From d8a1b678a4ff9e9eeb85a5cab87471a3203655d6 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Sun, 30 Aug 2026 19:10:45 -0500 Subject: [PATCH 1/7] Pin FLUX.2's revision, and give resident LTX its text encoder Two things stood between the resident API and a host serving it. `flux2_load_pipeline()` reached the hub at hfhub's default revision, the branch `main` -- resolved through `refs/main` and then over the network. A read-only bind of one snapshot carries neither, so the VAE, the Qwen3 encoder and the tokenizer all failed to resolve inside a container. `revision` now threads to all four `.flux2_cached` calls and a branch name is refused rather than passed through. `resident_load("ltx")` loaded a pipeline that could not generate. `ltx23_load_pipeline()` does not load the text encoder and `txt2vid_ltx2()` takes it per call, so the handle activated fine and the first generation had no encoder. It now loads Gemma3 pinned on the CPU and injects it, which also stops a serving caller re-reading 7.6 GB of encoder for every clip. The encoder is deliberately kept OUT of `staging`: that list is what activation moves to the device, and the encoder belongs on the host. Its pinned bytes are still counted, so the handle's declared host footprint is the real one. --- R/resident.R | 74 +++++++++++++++++++++++++++++++++++++- R/txt2img_flux2.R | 58 ++++++++++++++++++++++++++---- man/flux2_load_pipeline.Rd | 10 +++++- man/resident_load.Rd | 13 ++++++- 4 files changed, 145 insertions(+), 10 deletions(-) diff --git a/R/resident.R b/R/resident.R index 7a61b38..1f9d7ff 100644 --- a/R/resident.R +++ b/R/resident.R @@ -235,6 +235,14 @@ #' \code{checkpoint_path}; \code{sdxl} needs nothing (it defaults to the #' \code{\link{download_sdxl}} cache). #' @param verbose Print progress messages. +#' @param text_encoder,tokenizer \code{ltx} only: paths to the Gemma3 encoder +#' artifact and the tokenizer directory. \code{\link{ltx23_load_pipeline}} +#' does not load these -- \code{\link{txt2vid_ltx2}} takes them per call -- +#' so a handle built without them can be activated and cannot generate. +#' Given here they are loaded ONCE, pinned on the host, and passed to every +#' generate; the encoder rides to the card for its phase and back off, the +#' way the pipeline's own components do, so it is never resident beside the +#' transformer. #' #' @return A \code{diffuseR_resident} handle (an environment). Inspect it #' with \code{\link{resident_status}}; the fields of interest are the @@ -256,8 +264,19 @@ #' @export resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", "sdxl", "sd21"), - device = "cuda", ..., verbose = TRUE) { + device = "cuda", ..., verbose = TRUE, + text_encoder = NULL, tokenizer = NULL) { model <- match.arg(model) + ## NAMED ARGUMENTS RATHER THAN `...`, because `...` goes to the family + ## loader and `ltx23_load_pipeline` has no `...` of its own -- an unknown + ## argument there is an error, not a pass-through. Refused for the other + ## families for the same reason: silently ignoring them would leave a + ## caller believing a text encoder had been loaded. + if (!identical(model, "ltx") && + (!is.null(text_encoder) || !is.null(tokenizer))) { + stop("text_encoder/tokenizer apply to the ltx family only; ", + model, " loads its own", call. = FALSE) + } if (!torch::cuda_is_available()) { stop("resident_load() requires CUDA", call. = FALSE) } @@ -308,6 +327,48 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", res$gpu_components <- pipeline$gpu_components res$components <- names(.resident_components(pipeline)) res$pinned_bytes <- .resident_pinned_bytes(staging) + + ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING. + ## + ## `txt2vid_ltx2` takes `text_encoder` and `tokenizer` per call and + ## accepts a PATH, which it then loads -- so a serving caller that passed + ## paths would re-read 7.6 GB of Gemma3 on every request. Loading here + ## makes it once. + ## + ## Deliberately NOT added to `staging`: `resident_activate` puts + ## everything in staging on the card at once, and the encoder beside the + ## transformer does not fit. It carries its own staging attribute from + ## `pin = TRUE`, and `encode_with_gemma3` onloads it for the encode and + ## offloads on exit -- one GPU tenant per phase, the same discipline the + ## pipeline's own components follow. + if (identical(model, "ltx") && !is.null(text_encoder)) { + if (is.null(tokenizer)) { + stop("text_encoder needs a tokenizer: the encode takes both", + call. = FALSE) + } + if (verbose) message("Loading the Gemma3 text encoder (pinned)...") + res$text_encoder <- load_gemma3_text_encoder( + text_encoder, device = "cpu", pin = TRUE, verbose = verbose) + res$tokenizer <- gemma3_tokenizer(tokenizer) + ## Counted, so `resident_status()` reports what the process actually + ## holds. A pinned set omitted from the total reads as headroom that + ## is not there, and the fleet's admission arithmetic is downstream + ## of this number. + ## WRAPPED IN A LIST, AND THAT IS NOT COSMETIC. There are two + ## staging shapes in this package: a pipeline's is a list OF + ## COMPONENTS each holding a list of pairs, which is why + ## `.resident_pinned_bytes` loops twice; an encoder's + ## `attr(model, "staging")` is a FLAT list of pairs, which is why + ## `.staged_onload` loops once. Passing the flat one straight in + ## reads a pair's fields as pairs and dies on `pair$pinned$shape` + ## -- "object of type 'closure' is not subsettable", from inside a + ## worker, sixty seconds after the pin began. + te_staging <- attr(res$text_encoder, "staging") + if (!is.null(te_staging)) { + res$pinned_bytes <- res$pinned_bytes + + .resident_pinned_bytes(list(te_staging)) + } + } res$state <- "inactive" res$last_error <- NULL res$loaded_at <- Sys.time() @@ -762,6 +823,17 @@ resident_generate <- function(res, prompt, ...) { } }), want) } + ## LTX takes its text encoder per call and stores none, so a handle that + ## loaded one has to hand it over on every generate. An explicit argument + ## still wins -- this fills a gap rather than overriding a caller who + ## brought precomputed embeds or a different encoder. + if (identical(res$model, "ltx") && !is.null(res$text_encoder) && + is.null(args$text_encoder) && is.null(args$prompt_embeds)) { + args$text_encoder <- res$text_encoder + if (is.null(args$tokenizer)) { + args$tokenizer <- res$tokenizer + } + } args } diff --git a/R/txt2img_flux2.R b/R/txt2img_flux2.R index b52b915..b3ba0c6 100644 --- a/R/txt2img_flux2.R +++ b/R/txt2img_flux2.R @@ -11,13 +11,46 @@ #' @name txt2img_flux2 NULL +# The revision a hub lookup resolves against. +# +# NULL means hfhub's default, "main" -- a BRANCH, which it resolves through +# `refs/main` in the cache and, failing that, over the network. An exact 40-hex +# commit takes hfhub's fast path instead: straight to +# `snapshots//`, never consulting `refs/`. +# +# That difference is what lets a pipeline load from a cache holding ONLY the +# snapshot -- a read-only bind of one revision, with no `refs/` and no network, +# which is how a fleet node serves weights it did not bake. It is also the +# stronger guarantee generally: a branch moves, so a deployment pinned to +# yesterday's weights would silently start resolving today's. +# +# A branch name is REFUSED rather than passed through. Accepting one would hand +# hfhub a value it resolves the slow way, which is the behaviour this argument +# exists to avoid, and the caller would have no way to tell. +.flux2_rev <- function(revision) { + if (is.null(revision)) { + return(list()) + } + if (!is.character(revision) || length(revision) != 1L || is.na(revision) || + !grepl("^[0-9a-f]{40}$", revision)) { + stop("revision must be a single 40-character hex commit, not a branch ", + "name: a branch resolves through refs/ and defeats the point of ", + "pinning one", call. = FALSE) + } + list(revision = revision) +} + # Resolve a FLUX.2-klein support file from the HuggingFace cache -.flux2_cached <- function(file) { +.flux2_cached <- function(file, revision = NULL) { if (!requireNamespace("hfhub", quietly = TRUE)) { stop("The hfhub package is required to locate model files.") } + ## Outside the tryCatch below, which reports every error as a missing + ## download: a branch name is a caller mistake, not an absent file. + rev <- .flux2_rev(revision) tryCatch( - hfhub::hub_download(.flux2_repo, file, local_files_only = TRUE), + do.call(hfhub::hub_download, + c(list(.flux2_repo, file, local_files_only = TRUE), rev)), error = function(e) { stop("Missing ", file, " in the HuggingFace cache; ", "run download_flux2_klein() first.", call. = FALSE) @@ -47,6 +80,12 @@ NULL #' resolves via \code{options(diffuseR.pin_staging)} then the #' host-RAM-aware \code{\link{recommend}} decision. #' @param verbose Logical. +#' @param revision Optional exact 40-hex commit for the support files (VAE, +#' Qwen3 encoder, tokenizer) this pulls from the \code{black-forest-labs} +#' cache. With one they resolve straight out of +#' \code{snapshots//}, so the load works against a cache holding +#' only that snapshot -- no \code{refs/} entry and no network. The +#' transformer comes from \code{model_dir} and is unaffected. #' #' @return A \code{flux2_pipeline} list. #' @@ -55,7 +94,10 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4", "bf16"), text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, pin = NULL, - verbose = TRUE) { + verbose = TRUE, revision = NULL) { + ## Checked before any device work, where a configuration mistake is + ## cheapest to report. + .flux2_rev(revision) precision <- .flux_resolve_precision(match.arg(precision), file.path(tools::R_user_dir("diffuseR", "data"), "flux2-klein-4b-")) if (is.null(text_device)) { @@ -124,10 +166,11 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading FLUX.2 VAE decoder...") } - vae_config <- jsonlite::fromJSON(.flux2_cached("vae/config.json")) + vae_config <- jsonlite::fromJSON(.flux2_cached("vae/config.json", + revision)) pipe$vae_bn_eps <- vae_config$batch_norm_eps %||% 1e-4 pipe$decoder <- load_flux2_vae_decoder( - .flux2_cached("vae/diffusion_pytorch_model.safetensors"), + .flux2_cached("vae/diffusion_pytorch_model.safetensors", revision), latent_channels = as.integer(vae_config$latent_channels %||% 32L), verbose = verbose ) @@ -142,13 +185,14 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading Qwen3 text encoder...") } - te_dir <- dirname(.flux2_cached("text_encoder/config.json")) + te_dir <- dirname(.flux2_cached("text_encoder/config.json", revision)) pipe$text_encoder <- load_qwen3_text_encoder( te_dir, device = if (phase_offload) "cpu" else text_device, dtype = if (text_device == "cpu") "float32" else "bfloat16", verbose = verbose ) - pipe$tokenizer <- qwen_bpe_tokenizer(.flux2_cached("tokenizer/tokenizer.json")) + pipe$tokenizer <- qwen_bpe_tokenizer( + .flux2_cached("tokenizer/tokenizer.json", revision)) components <- c("transformer", "decoder") if (!identical(text_device, "cpu")) { diff --git a/man/flux2_load_pipeline.Rd b/man/flux2_load_pipeline.Rd index 0ea182a..809e978 100644 --- a/man/flux2_load_pipeline.Rd +++ b/man/flux2_load_pipeline.Rd @@ -11,7 +11,8 @@ flux2_load_pipeline( attn_chunk = NULL, phase_offload = TRUE, pin = NULL, - verbose = TRUE + verbose = TRUE, + revision = NULL ) } \arguments{ @@ -37,6 +38,13 @@ resolves via \code{options(diffuseR.pin_staging)} then the host-RAM-aware \code{\link{recommend}} decision.} \item{verbose}{Logical.} + +\item{revision}{Optional exact 40-hex commit for the support files (VAE, +Qwen3 encoder, tokenizer) this pulls from the \code{black-forest-labs} +cache. With one they resolve straight out of +\code{snapshots//}, so the load works against a cache holding +only that snapshot -- no \code{refs/} entry and no network. The +transformer comes from \code{model_dir} and is unaffected.} } \value{ A \code{flux2_pipeline} list. diff --git a/man/resident_load.Rd b/man/resident_load.Rd index 8f0bcd1..808e95c 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -7,7 +7,9 @@ resident_load( model = c("flux2", "flux1", "zimage", "ltx", "sdxl", "sd21"), device = "cuda", ..., - verbose = TRUE + verbose = TRUE, + text_encoder = NULL, + tokenizer = NULL ) } \arguments{ @@ -23,6 +25,15 @@ resident_load( \code{\link{download_sdxl}} cache).} \item{verbose}{Print progress messages.} + +\item{text_encoder,tokenizer}{\code{ltx} only: paths to the Gemma3 encoder +artifact and the tokenizer directory. \code{\link{ltx23_load_pipeline}} +does not load these -- \code{\link{txt2vid_ltx2}} takes them per call -- +so a handle built without them can be activated and cannot generate. +Given here they are loaded ONCE, pinned on the host, and passed to every +generate; the encoder rides to the card for its phase and back off, the +way the pipeline's own components do, so it is never resident beside the +transformer.} } \value{ A \code{diffuseR_resident} handle (an environment). Inspect it From 9b25c2e9ebe6341572f611f2db911126b4383cf3 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Sun, 30 Aug 2026 19:10:46 -0500 Subject: [PATCH 2/7] Bump version to 0.2.2.8 --- DESCRIPTION | 2 +- NEWS.md | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index c047ea6..eca4757 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.7 +Version: 0.2.2.8 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 04a8ae9..a046aa7 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,24 @@ +# diffuseR 0.2.2.8 + +* `flux2_load_pipeline()` takes a `revision`. Its VAE, Qwen3 encoder and + tokenizer come from the Hugging Face cache, and hfhub's default revision + is the branch `main` -- resolved through `refs/main` and, failing that, + over the network. A read-only bind of one snapshot carries neither, so + the load failed there; an exact 40-hex commit takes hfhub straight to + `snapshots//`. A branch name is refused rather than + passed through. + +* `resident_load("ltx", ...)` takes `text_encoder` and `tokenizer` paths. + `ltx23_load_pipeline()` does not load them and `txt2vid_ltx2()` takes + them per call, so a resident LTX handle could be activated and could not + generate -- and a serving caller passing paths re-read 7.6 GB of Gemma3 + on every request. Given here they load once, pinned on the host, and + `resident_generate()` supplies them. They stay OUT of the handle's + staging on purpose: `resident_activate()` places everything in staging at + once, and the encoder does not fit beside the transformer. It rides to + the card for its own phase and back off, as the pipeline's components do. + `pinned_bytes` counts it. + # diffuseR 0.2.2.7 * `recommend()` diagnosed the wrong safetensors capability for bf16. The From bb923bb195e5377ca2f964245fb2dfecfe57a1d6 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Tue, 1 Sep 2026 17:10:22 -0500 Subject: [PATCH 3/7] Stage the resident LTX prompt encode to the card instead of CPU txt2vid_ltx2() chose the encode device with `if (is.character(text_encoder)) device else "cpu"`, so a preloaded encoder -- the resident/gpuhost path, resident_load("ltx", text_encoder = ...) -- always encoded the prompt on CPU. That path loads the encoder with pin = TRUE precisely so encode_with_gemma3() can DMA it onto the card per encode and free it after, but the `else "cpu"` forced CPU and left the pinned staging unused: every prompt paid the ~24 s CPU encode instead of the ~7 s staged-GPU one, and on the gpuhost path that is once per chunk. The decision moves to .ltx23_text_encode_device(): a path loads onto the asked-for device; a preloaded encoder that carries a `staging` set and got a cuda request stages to the card; a bare CPU-resident object still degrades to CPU, because a cuda request without staging would send the tokens to the card while the weights sat on the host. Split out so the rule is asserted without a GPU or a real encode (test_text_encode_device.R, 8 cases), the way .resident_gen_args is. The in-process backend is unaffected: it precomputes connector_embeds and never reaches this branch. --- R/txt2vid_ltx23.R | 28 +++++++++++++++++++++- inst/tinytest/test_text_encode_device.R | 31 +++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 inst/tinytest/test_text_encode_device.R diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 8501d7d..89adbce 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -373,6 +373,32 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", structure(pipe, class = "ltx23_pipeline") } +# The device the prompt encode runs on, for the text_encoder txt2vid_ltx2 +# was handed. A PATH is loaded fresh onto `device` -- a caller who names a +# file is asking for it there. A PRELOADED object is CPU-resident by the +# resident loader's convention (load_gemma3_text_encoder(device = "cpu", +# pin = TRUE)), and only safe to encode on the compute device when it +# carries the pinned `staging` set that lets encode_with_gemma3() DMA it +# there and back per encode. Without that set, `device = "cuda"` would send +# the tokens to the card while the weights sat on the host -- a mismatch -- +# so a bare object degrades to CPU. +# +# The `else "cpu"` this replaces was blunter than that: it forced CPU for +# EVERY preloaded encoder, so the resident/gpuhost path (which loads with +# pin = TRUE precisely to stage) encoded every prompt on CPU (~24 s) with +# its own pinned staging sitting unused, instead of the ~7 s staged-GPU +# encode. Split out so the decision is assertable without a GPU or a real +# encode, the way `.resident_gen_args` is. +.ltx23_text_encode_device <- function(text_encoder, device) { + if (is.character(text_encoder)) { + return(device) + } + if (!is.null(attr(text_encoder, "staging")) && grepl("^cuda", device)) { + return(device) + } + "cpu" +} + #' Generate video (and audio) with LTX-2.3 #' #' Distilled text-to-video generation: encodes the prompt with Gemma3 + @@ -561,7 +587,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, prompt, model = text_encoder, tokenizer = tokenizer, max_sequence_length = max_sequence_length, - device = if (is.character(text_encoder)) device else "cpu", + device = .ltx23_text_encode_device(text_encoder, device), verbose = verbose ) } diff --git a/inst/tinytest/test_text_encode_device.R b/inst/tinytest/test_text_encode_device.R new file mode 100644 index 0000000..44382f7 --- /dev/null +++ b/inst/tinytest/test_text_encode_device.R @@ -0,0 +1,31 @@ +# .ltx23_text_encode_device: which device the prompt encode runs on for +# the text_encoder txt2vid_ltx2 was handed. Pure decision, no torch, no +# GPU, no real encode -- the point is to pin the rule that a resident +# (preloaded, pinned) encoder stages to the card instead of encoding on +# CPU, which is what the gpuhost/resident path silently did before. + +library(diffuseR) +f <- diffuseR:::.ltx23_text_encode_device + +# A PATH loads fresh onto whatever device was asked for. +expect_equal(f("/models/gemma3-nf4", "cuda"), "cuda") +expect_equal(f("/models/gemma3-nf4", "cuda:0"), "cuda:0") +expect_equal(f("/models/gemma3-nf4", "cpu"), "cpu") + +# A PRELOADED object is only encoded on the card when it carries the +# pinned staging set (the resident loader's pin = TRUE). This is the case +# that regressed: the resident/gpuhost encoder HAS staging and a cuda +# request, so it must stage to the card, not fall to CPU. +staged <- structure(list(), staging = list(TRUE)) +expect_equal(f(staged, "cuda"), "cuda") +expect_equal(f(staged, "cuda:0"), "cuda:0") +# An explicit cpu request is still honoured even with staging present. +expect_equal(f(staged, "cpu"), "cpu") + +# A preloaded object with NO staging cannot run on the card -- its weights +# sit on the host and a cuda request would be a device mismatch -- so it +# degrades to CPU. This is the safety the blunt `else "cpu"` was protecting, +# preserved for exactly this case. +bare <- structure(list()) +expect_equal(f(bare, "cuda"), "cpu") +expect_equal(f(bare, "cpu"), "cpu") From 4a53899e817496dd9875fabc58bdb426c885b456 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Tue, 1 Sep 2026 17:10:22 -0500 Subject: [PATCH 4/7] Bump version to 0.2.2.9 --- DESCRIPTION | 2 +- NEWS.md | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index eca4757..481819f 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.8 +Version: 0.2.2.9 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 a046aa7..003c6a8 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,20 @@ +# diffuseR 0.2.2.9 + +* **A resident LTX encoder now stages the prompt encode to the card + instead of running it on CPU.** `txt2vid_ltx2()` chose the encode device + with `if (is.character(text_encoder)) device else "cpu"` -- so a + PRELOADED encoder (the resident/gpuhost path, `resident_load("ltx", + text_encoder = ...)`) always encoded on CPU, even though the resident + loader page-locks it with `pin = TRUE` for exactly the staged transfer + `encode_with_gemma3()` supports. The pinned staging sat unused and every + prompt paid the ~24 s CPU encode instead of the ~7 s staged-GPU one; on + the gpuhost path that is once per chunk. The device decision is now + `.ltx23_text_encode_device()`: a path loads onto the asked-for device, a + preloaded encoder with a `staging` set and a cuda request stages to the + card, and a bare CPU-resident object still degrades to CPU (a cuda + request without staging would be a device mismatch). Pure and unit-tested + without a GPU (`test_text_encode_device.R`). + # diffuseR 0.2.2.8 * `flux2_load_pipeline()` takes a `revision`. Its VAE, Qwen3 encoder and From e1777b17b8097802cee2123f5868215aa08cec6d Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 11 Sep 2026 05:55:26 -0500 Subject: [PATCH 5/7] Complete a partial staged onload instead of calling it resident The Gemma3 encoder's staged encode and the LTX pipeline's per-phase onload both decided "already on the card" by probing the FIRST staging pair. An onload that dies partway -- device memory runs out with most of the encoder already copied -- leaves exactly that pair on the card and the rest on the host, so every later call skipped the onload and failed on the first matrix multiply with "mat2 is on cpu", on every request, until the process ended. That is what took the gpuhost's ltx-2.3 entry down for USA 20260912 on 2026-09-10: one failed encoder onload at track 2, then the same refusal from every request after it, with nothing left in the process that could move the encoder back. Three changes. `.staged_on()` asks every pair, not the first one. `.staged_onload()` is idempotent per pair: a resident tensor is left where it is (no re-transfer over itself, which is what the first-pair probe existed to avoid) and a half-copied component is completed rather than restarted. And `encode_with_gemma3()` arms its offload BEFORE the onload, so a failed transfer is undone on the way out and the next encode starts from a clean host copy and reports the real error. The LTX onload closure drops its own probe and lets the per-pair rule decide; plain modules without staging keep the `$parameters` probe they had. Covered by test_staged_on.R (pure fakes, no torch) and a partial round trip added to test_staging.R (CUDA). --- R/gemma3_text_encoder.R | 24 +++++++---- R/staging.R | 58 +++++++++++++++++++++++++ R/txt2vid_ltx23.R | 41 ++++++++---------- inst/tinytest/test_staged_on.R | 77 ++++++++++++++++++++++++++++++++++ inst/tinytest/test_staging.R | 17 ++++++++ man/dot-staged_on.Rd | 33 +++++++++++++++ man/dot-staged_onload.Rd | 8 ++++ 7 files changed, 227 insertions(+), 31 deletions(-) create mode 100644 inst/tinytest/test_staged_on.R create mode 100644 man/dot-staged_on.Rd diff --git a/R/gemma3_text_encoder.R b/R/gemma3_text_encoder.R index 4ac4de5..0d977f3 100644 --- a/R/gemma3_text_encoder.R +++ b/R/gemma3_text_encoder.R @@ -860,14 +860,24 @@ encode_with_gemma3 <- function(prompts, model = NULL, tokenizer = NULL, } # A pinned CPU-resident model (see the loaders' pin argument) swaps - # to the compute device for the encode and back for free afterwards + # to the compute device for the encode and back for free afterwards. + # A model already on the device IN FULL is left there, and left alone + # on exit: whoever put it there owns it. + # + # THE OFFLOAD IS ARMED BEFORE THE ONLOAD, not after it. An onload that + # fails partway -- device memory runs out with most of the encoder + # already copied -- must still be undone on the way out, or the + # encoder stays half on the card with nothing left to move it back, + # and every encode after it inherits the split. `.staged_on` asks + # every pair rather than the first one for the same reason: the first + # tensor of a half-copied encoder IS on the card, and a probe of it + # alone reported the encoder resident and skipped the onload on every + # request for the rest of the process (the gpuhost's ltx-2.3 entry, + # 2026-09-10: "mat2 is on cpu" from every encode after one failure). staging <- attr(model, "staging") - if (!is.null(staging) && device != "cpu") { - cur <- tryCatch(staging[[1]]$live$device$type, error = function(e) NULL) - if (!identical(cur, device)) { - .staged_onload(staging, device) - on.exit(.staged_offload(staging), add = TRUE) - } + if (!is.null(staging) && device != "cpu" && !.staged_on(staging, device)) { + on.exit(.staged_offload(staging), add = TRUE) + .staged_onload(staging, device) } # Ensure prompts is a list diff --git a/R/staging.R b/R/staging.R index b5b5cea..19d6114 100644 --- a/R/staging.R +++ b/R/staging.R @@ -80,14 +80,72 @@ NULL }, error = function(e) NULL) } +# The device TYPE a device spec names: "cuda" for "cuda", "cuda:0" and a +# torch_device on the card alike. Type is what the staging checks +# compare -- a pair lives on the card or on the host -- and a string +# compare keeps `.staged_on` runnable without torch, which is how it is +# tested. +.device_type <- function(device) { + if (inherits(device, "torch_device")) { + return(device$type) + } + sub(":.*$", "", as.character(device)) +} + +#' Is every pinned tensor of a component on this device? +#' +#' The check a caller makes before skipping an onload. It asks EVERY +#' pair, not the first one: a component is on the card when all of it +#' is, and a probe of one tensor cannot tell a resident component from +#' one whose onload failed partway. That partial state is real -- an +#' onload that runs out of device memory leaves the pairs it copied on +#' the card and the rest on the host -- and a first-pair probe reports +#' it as "already resident", so every later phase skips the onload and +#' dies on a device mismatch, on every call, until the process ends. +#' That is how a gpuhost's LTX entry wedged for a whole show on +#' 2026-09-10: one failed encoder onload, then "mat2 is on cpu" from +#' every request after it. +#' +#' @param staging A component's staging set: the list of +#' \code{list(live, pinned)} pairs \code{.pin_component} returned. +#' @param device The compute device, as a string (\code{"cuda"}, +#' \code{"cuda:0"}) or a \code{torch_device}; only its type is compared. +#' @return TRUE when every pair's live tensor is on the device's type; +#' FALSE on any mismatch or unreadable pair. Vacuously TRUE for an +#' empty staging set, which holds nothing to move. +#' @keywords internal +.staged_on <- function(staging, device) { + type <- .device_type(device) + for (pair in staging) { + cur <- tryCatch(pair$live$device$type, + error = function(e) NA_character_) + if (!identical(cur, type)) { + return(FALSE) + } + } + TRUE +} + #' Move a pinned component onto the compute device #' #' Non-blocking copies from pinned memory share the default stream, #' so later kernels are ordered after them; no explicit sync needed. #' +#' Idempotent PER PAIR: a tensor already on the device is left where it +#' is, so a resident component costs nothing to onload again (no +#' re-transfer of weights over themselves, which fragments the +#' allocator pool) and a component whose earlier onload stopped partway +#' is completed rather than restarted. +#' #' @keywords internal .staged_onload <- function(staging, device) { + type <- .device_type(device) for (pair in staging) { + cur <- tryCatch(pair$live$device$type, + error = function(e) NA_character_) + if (identical(cur, type)) { + next + } pair$live$set_data(pair$pinned$to(device = device, non_blocking = TRUE)) } invisible(NULL) diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 89adbce..7767cc4 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -608,37 +608,30 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, } if (phase_offload) { # Idempotent: a resident component is already in place on - # the second and later calls of a chained run. Probe the - # staging pair's live tensor when staging exists - custom - # module classes (the NF4 transformer) may not expose + # the second and later calls of a chained run, and must not + # be re-transferred over itself (that fragments the + # allocator pool until the next large allocation OOMs). + # With staging, `.staged_onload` decides PER PAIR: a + # resident component is a no-op, and one whose earlier + # onload stopped partway is completed instead of being + # reported resident by its first tensor and left split. + # Without staging, probe the module -- custom module + # classes (the NF4 transformer) may not expose # $parameters, and a failed probe must not degrade into a - # re-onload: re-transferring resident weights over - # themselves fragments the allocator pool until the next - # large allocation OOMs. - if (is.character(what)) { - st_probe <- staging[[what]] - } else { - st_probe <- NULL - } - cur <- tryCatch({ - if (!is.null(st_probe)) { - st_probe[[1]]$live$device$type - } else { - module$parameters[[1]]$device$type - } - }, error = function(e) NULL) - if (identical(cur, target_type)) { - return(module) - } + # re-onload. if (is.character(what)) { st <- staging[[what]] } else { st <- NULL } - if (is.null(st)) { - module$to(device = device) - } else { + if (!is.null(st)) { .staged_onload(st, device) + return(module) + } + cur <- tryCatch(module$parameters[[1]]$device$type, + error = function(e) NULL) + if (!identical(cur, target_type)) { + module$to(device = device) } } module diff --git a/inst/tinytest/test_staged_on.R b/inst/tinytest/test_staged_on.R new file mode 100644 index 0000000..d8d7890 --- /dev/null +++ b/inst/tinytest/test_staged_on.R @@ -0,0 +1,77 @@ +# .staged_on / .staged_onload (R/staging.R): the checks a phase makes +# before moving a pinned component. Pure fakes, no torch, no GPU: a +# "pair" is anything with $live$device$type, $live$set_data and +# $pinned$to, which is all the helpers touch. +# +# The case that matters is the PARTIAL one. An onload that dies partway +# leaves the first pairs on the card and the rest on the host. Probing +# the first pair alone called that resident, so every later phase +# skipped the onload and failed on a device mismatch until the process +# ended (the gpuhost's ltx-2.3 entry, 2026-09-10). + +library(diffuseR) +staged_on <- diffuseR:::.staged_on +staged_onload <- diffuseR:::.staged_onload +device_type <- diffuseR:::.device_type + +# A fake pair whose live tensor sits on `type`, recording every +# set_data it receives. +fake_pair <- function(type) { + log <- new.env() + log$calls <- list() + live <- list(device = list(type = type), + set_data = function(x) { + log$calls[[length(log$calls) + 1L]] <- x + }) + pinned <- list(to = function(device, non_blocking = FALSE) { + paste0("copy->", device) + }) + list(live = live, pinned = pinned, log = log) +} +calls <- function(p) p$log$calls + +# Device type: strings with and without an index, and a torch_device. +expect_equal(device_type("cuda"), "cuda") +expect_equal(device_type("cuda:0"), "cuda") +expect_equal(device_type("cpu"), "cpu") +expect_equal(device_type(structure(list(type = "cuda", index = 0L), + class = "torch_device")), "cuda") + +# All on the card: resident. +st <- list(fake_pair("cuda"), fake_pair("cuda"), fake_pair("cuda")) +expect_true(staged_on(st, "cuda")) +expect_true(staged_on(st, "cuda:0")) +expect_false(staged_on(st, "cpu")) + +# All on the host: not resident. +st <- list(fake_pair("cpu"), fake_pair("cpu")) +expect_false(staged_on(st, "cuda")) +expect_true(staged_on(st, "cpu")) + +# PARTIAL: first pair on the card, second on the host. The first-pair +# probe said "resident"; the whole-set check must not. +st <- list(fake_pair("cuda"), fake_pair("cpu")) +expect_false(staged_on(st, "cuda")) + +# An unreadable pair is not resident. +st <- list(fake_pair("cuda"), list(live = NULL, pinned = NULL)) +expect_false(staged_on(st, "cuda")) + +# Empty staging holds nothing to move. +expect_true(staged_on(list(), "cuda")) + +# Onload is idempotent per pair: resident pairs are untouched, the rest +# are copied. A partial onload is completed, not restarted. +st <- list(fake_pair("cuda"), fake_pair("cpu"), + fake_pair("cuda"), fake_pair("cpu")) +staged_onload(st, "cuda") +expect_equal(length(calls(st[[1]])), 0L) +expect_equal(calls(st[[2]]), list("copy->cuda")) +expect_equal(length(calls(st[[3]])), 0L) +expect_equal(calls(st[[4]]), list("copy->cuda")) + +# A fully resident set is a no-op, whichever spelling names the card. +st <- list(fake_pair("cuda"), fake_pair("cuda")) +staged_onload(st, "cuda:0") +expect_equal(length(calls(st[[1]])), 0L) +expect_equal(length(calls(st[[2]])), 0L) diff --git a/inst/tinytest/test_staging.R b/inst/tinytest/test_staging.R index 0490091..0c64e4f 100644 --- a/inst/tinytest/test_staging.R +++ b/inst/tinytest/test_staging.R @@ -43,3 +43,20 @@ diffuseR:::.staged_onload(st, "cuda") torch::with_no_grad(out_gpu2 <- m(x$to(device = "cuda"))$cpu()) expect_true(as.numeric((out_gpu2 - out_gpu)$abs()$max()) == 0) diffuseR:::.staged_offload(st) + +# PARTIAL onload, the state a failed transfer leaves behind: the first +# pair on the card, the rest on the host. The whole-set check must call +# it not resident, and the next onload must complete it -- moving only +# what is missing -- rather than skip it on the strength of pair 1. +st[[1]]$live$set_data(st[[1]]$pinned$to(device = "cuda")) +expect_equal(st[[1]]$live$device$type, "cuda") +expect_equal(st[[length(st)]]$live$device$type, "cpu") +expect_false(diffuseR:::.staged_on(st, "cuda")) +diffuseR:::.staged_onload(st, "cuda") +expect_true(diffuseR:::.staged_on(st, "cuda")) +torch::with_no_grad(out_gpu3 <- m(x$to(device = "cuda"))$cpu()) +expect_true(as.numeric((out_gpu3 - out_gpu)$abs()$max()) == 0) +diffuseR:::.staged_offload(st) +expect_true(diffuseR:::.staged_on(st, "cpu")) +torch::with_no_grad(out_back2 <- m(x)) +expect_true(as.numeric((out_back2 - ref)$abs()$max()) == 0) diff --git a/man/dot-staged_on.Rd b/man/dot-staged_on.Rd new file mode 100644 index 0000000..0e91cf3 --- /dev/null +++ b/man/dot-staged_on.Rd @@ -0,0 +1,33 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.staged_on} +\alias{.staged_on} +\title{Is every pinned tensor of a component on this device?} +\usage{ +.staged_on(staging, device) +} +\arguments{ +\item{staging}{A component's staging set: the list of +\code{list(live, pinned)} pairs \code{.pin_component} returned.} + +\item{device}{The compute device, as a string (\code{"cuda"}, +\code{"cuda:0"}) or a \code{torch_device}; only its type is compared.} +} +\value{ +TRUE when every pair's live tensor is on the device's type; + FALSE on any mismatch or unreadable pair. Vacuously TRUE for an + empty staging set, which holds nothing to move. +} +\description{ +The check a caller makes before skipping an onload. It asks EVERY +pair, not the first one: a component is on the card when all of it +is, and a probe of one tensor cannot tell a resident component from +one whose onload failed partway. That partial state is real -- an +onload that runs out of device memory leaves the pairs it copied on +the card and the rest on the host -- and a first-pair probe reports +it as "already resident", so every later phase skips the onload and +dies on a device mismatch, on every call, until the process ends. +That is how a gpuhost's LTX entry wedged for a whole show on +2026-09-10: one failed encoder onload, then "mat2 is on cpu" from +every request after it. +} +\keyword{internal} diff --git a/man/dot-staged_onload.Rd b/man/dot-staged_onload.Rd index b9bb780..ace3459 100644 --- a/man/dot-staged_onload.Rd +++ b/man/dot-staged_onload.Rd @@ -8,5 +8,13 @@ \description{ Non-blocking copies from pinned memory share the default stream, so later kernels are ordered after them; no explicit sync needed. +} +\details{ +Idempotent PER PAIR: a tensor already on the device is left where it +is, so a resident component costs nothing to onload again (no +re-transfer of weights over themselves, which fragments the +allocator pool) and a component whose earlier onload stopped partway +is completed rather than restarted. + } \keyword{internal} From ded28f48ff75f5fac37fc534c09245ec1e2ae182 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 11 Sep 2026 05:55:26 -0500 Subject: [PATCH 6/7] Bump version to 0.2.2.10 --- DESCRIPTION | 2 +- NEWS.md | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index 481819f..53f5639 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.9 +Version: 0.2.2.10 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 003c6a8..d53af46 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,22 @@ +# diffuseR 0.2.2.10 + +* **A pinned component whose onload failed partway no longer stays + wedged.** The Gemma3 encoder's staged encode, and the LTX pipeline's + per-phase onload, decided "already on the card" by probing the FIRST + staging pair. An onload that dies partway -- device memory runs out + with most of the encoder copied -- leaves exactly that pair on the card + and the rest on the host, so every later call skipped the onload and + failed on the first matrix multiply with "mat2 is on cpu", on every + request, until the process ended. That is what took the gpuhost's + ltx-2.3 entry down for USA 20260912 on 2026-09-10. Three changes: + `.staged_on()` asks every pair, not the first; `.staged_onload()` is + idempotent per pair, so a resident component is a no-op (no re-transfer + over itself) and a partial one is completed; and `encode_with_gemma3()` + arms its offload BEFORE the onload, so a failed transfer is undone on + the way out and the next encode starts from a clean host copy. Covered + by `test_staged_on.R` (pure fakes, no GPU) and a partial round trip in + `test_staging.R` (CUDA). + # diffuseR 0.2.2.9 * **A resident LTX encoder now stages the prompt encode to the card From b8ab65487f2dab328dd24244705f080afc49500b Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 11 Sep 2026 07:12:04 -0500 Subject: [PATCH 7/7] Compare the card in staging checks, and drop the LTX encoder on unload Two findings from the 2026-09-11 review of the staging fix. `.staged_on()` and `.staged_onload()` compared device TYPE only, so a request for "cuda:1" counted a tensor on cuda:0 as resident, skipped the transfer, and left the weights on the wrong card. The old first-pair probe had the same blind spot; the new helper just made it explicit. `.device_spec()` now carries the index when the caller named one, and `.on_device()` requires it to match. Bare "cuda" still accepts any card. `resident_unload()` cleared `staging` and set pinned_bytes to 0 but never touched `res$text_encoder`, which resident_load("ltx") keeps outside `staging` on purpose. An unloaded handle therefore held the Gemma3 encoder's pinned buffers while reporting none. Unload now offloads the encoder's staging and drops the encoder and tokenizer. Both covered without a GPU: the card cases in test_staged_on.R and an unload-with-encoder case in test_resident.R. --- NEWS.md | 12 +++++++ R/resident.R | 9 +++++ R/staging.R | 60 ++++++++++++++++++++++------------ inst/tinytest/test_resident.R | 21 ++++++++++++ inst/tinytest/test_staged_on.R | 44 +++++++++++++++++-------- man/dot-staged_on.Rd | 9 ++--- 6 files changed, 118 insertions(+), 37 deletions(-) diff --git a/NEWS.md b/NEWS.md index d53af46..1750bc5 100644 --- a/NEWS.md +++ b/NEWS.md @@ -17,6 +17,18 @@ by `test_staged_on.R` (pure fakes, no GPU) and a partial round trip in `test_staging.R` (CUDA). +* **Staging compares the card, not just the device type.** A request + for `"cuda:1"` no longer counts a tensor on `cuda:0` as resident, so a + multi-GPU caller asking for a particular card gets its weights moved + there instead of a skipped transfer and a device mismatch. A request + for bare `"cuda"` still accepts any card, as before. + +* **`resident_unload()` drops the LTX text encoder.** The encoder a + resident LTX handle loads (0.2.2.8) sits outside `staging` by design, + and unload never released it: an unloaded handle kept the encoder's + pinned buffers while reporting `pinned_bytes = 0`. Both review findings + from the 2026-09-11 Codex pass. + # diffuseR 0.2.2.9 * **A resident LTX encoder now stages the prompt encode to the card diff --git a/R/resident.R b/R/resident.R index 1f9d7ff..0ace6d7 100644 --- a/R/resident.R +++ b/R/resident.R @@ -903,8 +903,17 @@ resident_unload <- function(res) { for (nm in names(res$staging)) { .staged_offload(res$staging[[nm]]) } + # The LTX text encoder is deliberately outside `staging` (see + # resident_load), but it is pinned host memory this handle owns + # and counted in pinned_bytes, so it goes the same way. + te_staging <- attr(res$text_encoder, "staging") + if (!is.null(te_staging)) { + .staged_offload(te_staging) + } }, error = function(e) NULL) res$pipeline <- NULL + res$text_encoder <- NULL + res$tokenizer <- NULL res$staging <- list() res$components <- character(0) res$pinned_bytes <- 0 diff --git a/R/staging.R b/R/staging.R index 19d6114..8705ad4 100644 --- a/R/staging.R +++ b/R/staging.R @@ -80,16 +80,39 @@ NULL }, error = function(e) NULL) } -# The device TYPE a device spec names: "cuda" for "cuda", "cuda:0" and a -# torch_device on the card alike. Type is what the staging checks -# compare -- a pair lives on the card or on the host -- and a string -# compare keeps `.staged_on` runnable without torch, which is how it is -# tested. -.device_type <- function(device) { +# What a device spec names, as type and index. The index is NA when the +# spec leaves it open: "cuda" means whichever card is current, so a +# target without an index accepts any card, and "cuda:1" accepts only +# that one. A tensor's own device always carries a concrete index on the +# card (torch reports 0 for "cuda"), so the comparison below is exact +# whenever the caller asked for a particular card. Parsed from the string +# rather than through torch_device() so `.staged_on` runs without torch, +# which is how it is tested. +.device_spec <- function(device) { if (inherits(device, "torch_device")) { - return(device$type) + return(list(type = device$type, + index = as.integer(device$index %||% NA_integer_))) } - sub(":.*$", "", as.character(device)) + s <- as.character(device) + index <- if (grepl(":", s, fixed = TRUE)) { + as.integer(sub("^[^:]*:", "", s)) + } else { + NA_integer_ + } + list(type = sub(":.*$", "", s), index = index) +} + +# Is this tensor on the device the spec names? Type must match; the +# index must match too when the spec has one. +.on_device <- function(tensor, spec) { + d <- tryCatch(tensor$device, error = function(e) NULL) + if (is.null(d) || !identical(d$type, spec$type)) { + return(FALSE) + } + if (is.na(spec$index)) { + return(TRUE) + } + identical(as.integer(d$index %||% NA_integer_), spec$index) } #' Is every pinned tensor of a component on this device? @@ -109,17 +132,16 @@ NULL #' @param staging A component's staging set: the list of #' \code{list(live, pinned)} pairs \code{.pin_component} returned. #' @param device The compute device, as a string (\code{"cuda"}, -#' \code{"cuda:0"}) or a \code{torch_device}; only its type is compared. -#' @return TRUE when every pair's live tensor is on the device's type; -#' FALSE on any mismatch or unreadable pair. Vacuously TRUE for an -#' empty staging set, which holds nothing to move. +#' \code{"cuda:1"}) or a \code{torch_device}. A spec without an index +#' accepts any card; one with an index accepts only that card. +#' @return TRUE when every pair's live tensor is on that device; FALSE +#' on any mismatch or unreadable pair. Vacuously TRUE for an empty +#' staging set, which holds nothing to move. #' @keywords internal .staged_on <- function(staging, device) { - type <- .device_type(device) + spec <- .device_spec(device) for (pair in staging) { - cur <- tryCatch(pair$live$device$type, - error = function(e) NA_character_) - if (!identical(cur, type)) { + if (!.on_device(pair$live, spec)) { return(FALSE) } } @@ -139,11 +161,9 @@ NULL #' #' @keywords internal .staged_onload <- function(staging, device) { - type <- .device_type(device) + spec <- .device_spec(device) for (pair in staging) { - cur <- tryCatch(pair$live$device$type, - error = function(e) NA_character_) - if (identical(cur, type)) { + if (.on_device(pair$live, spec)) { next } pair$live$set_data(pair$pinned$to(device = device, non_blocking = TRUE)) diff --git a/inst/tinytest/test_resident.R b/inst/tinytest/test_resident.R index c5024ca..ec74cd6 100644 --- a/inst/tinytest/test_resident.R +++ b/inst/tinytest/test_resident.R @@ -166,6 +166,27 @@ expect_error(resident_activate(u), pattern = "unloaded") # Status still works on an unloaded handle. expect_equal(resident_status(u)$state, "unloaded") +# Unload drops the LTX text encoder too. It lives outside `staging` on +# purpose (resident_load), so a handle that only cleared `staging` kept +# the encoder's pinned buffers while reporting pinned_bytes = 0. +# A fake pair records the offload it receives. +te_log <- new.env() +te_log$calls <- list() +te_pair <- list(live = list(set_data = function(x) { + te_log$calls[[length(te_log$calls) + 1L]] <- x + }), + pinned = "pinned-copy") +u2 <- mk("inactive") +u2$text_encoder <- structure(list(), staging = list(te_pair)) +u2$tokenizer <- list(vocab = 3L) +resident_unload(u2) +expect_equal(u2$state, "unloaded") +expect_null(u2$text_encoder) +expect_null(u2$tokenizer) +# The encoder's pinned staging was offloaded (pointer swap back to the +# pinned copy) before being dropped. +expect_equal(te_log$calls, list("pinned-copy")) + # --- CUDA round trip -------------------------------------------------------------- if (have_torch && at_home() && torch::cuda_is_available()) { diff --git a/inst/tinytest/test_staged_on.R b/inst/tinytest/test_staged_on.R index d8d7890..710a5cf 100644 --- a/inst/tinytest/test_staged_on.R +++ b/inst/tinytest/test_staged_on.R @@ -1,7 +1,7 @@ # .staged_on / .staged_onload (R/staging.R): the checks a phase makes # before moving a pinned component. Pure fakes, no torch, no GPU: a -# "pair" is anything with $live$device$type, $live$set_data and -# $pinned$to, which is all the helpers touch. +# "pair" is anything with $live$device, $live$set_data and $pinned$to, +# which is all the helpers touch. # # The case that matters is the PARTIAL one. An onload that dies partway # leaves the first pairs on the card and the rest on the host. Probing @@ -12,14 +12,14 @@ library(diffuseR) staged_on <- diffuseR:::.staged_on staged_onload <- diffuseR:::.staged_onload -device_type <- diffuseR:::.device_type +device_spec <- diffuseR:::.device_spec -# A fake pair whose live tensor sits on `type`, recording every -# set_data it receives. -fake_pair <- function(type) { +# A fake pair whose live tensor sits on `type` (and card `index`), +# recording every set_data it receives. +fake_pair <- function(type, index = if (type == "cpu") NULL else 0L) { log <- new.env() log$calls <- list() - live <- list(device = list(type = type), + live <- list(device = list(type = type, index = index), set_data = function(x) { log$calls[[length(log$calls) + 1L]] <- x }) @@ -30,12 +30,16 @@ fake_pair <- function(type) { } calls <- function(p) p$log$calls -# Device type: strings with and without an index, and a torch_device. -expect_equal(device_type("cuda"), "cuda") -expect_equal(device_type("cuda:0"), "cuda") -expect_equal(device_type("cpu"), "cpu") -expect_equal(device_type(structure(list(type = "cuda", index = 0L), - class = "torch_device")), "cuda") +# Device spec: strings with and without an index, and a torch_device. +expect_equal(device_spec("cuda"), list(type = "cuda", index = NA_integer_)) +expect_equal(device_spec("cuda:1"), list(type = "cuda", index = 1L)) +expect_equal(device_spec("cpu"), list(type = "cpu", index = NA_integer_)) +expect_equal(device_spec(structure(list(type = "cuda", index = 1), + class = "torch_device")), + list(type = "cuda", index = 1L)) +expect_equal(device_spec(structure(list(type = "cuda", index = NULL), + class = "torch_device")), + list(type = "cuda", index = NA_integer_)) # All on the card: resident. st <- list(fake_pair("cuda"), fake_pair("cuda"), fake_pair("cuda")) @@ -53,6 +57,14 @@ expect_true(staged_on(st, "cpu")) st <- list(fake_pair("cuda"), fake_pair("cpu")) expect_false(staged_on(st, "cuda")) +# THE WRONG CARD is not this card. A request naming an index accepts only +# that index; a request without one accepts any card. +st <- list(fake_pair("cuda", 0L), fake_pair("cuda", 1L)) +expect_false(staged_on(st, "cuda:0")) +expect_false(staged_on(st, "cuda:1")) +expect_true(staged_on(st, "cuda")) +expect_true(staged_on(list(fake_pair("cuda", 1L)), "cuda:1")) + # An unreadable pair is not resident. st <- list(fake_pair("cuda"), list(live = NULL, pinned = NULL)) expect_false(staged_on(st, "cuda")) @@ -75,3 +87,9 @@ st <- list(fake_pair("cuda"), fake_pair("cuda")) staged_onload(st, "cuda:0") expect_equal(length(calls(st[[1]])), 0L) expect_equal(length(calls(st[[2]])), 0L) + +# A pair on another card IS moved when a particular card was asked for. +st <- list(fake_pair("cuda", 0L), fake_pair("cuda", 1L)) +staged_onload(st, "cuda:1") +expect_equal(calls(st[[1]]), list("copy->cuda:1")) +expect_equal(length(calls(st[[2]])), 0L) diff --git a/man/dot-staged_on.Rd b/man/dot-staged_on.Rd index 0e91cf3..90171e6 100644 --- a/man/dot-staged_on.Rd +++ b/man/dot-staged_on.Rd @@ -10,12 +10,13 @@ \code{list(live, pinned)} pairs \code{.pin_component} returned.} \item{device}{The compute device, as a string (\code{"cuda"}, -\code{"cuda:0"}) or a \code{torch_device}; only its type is compared.} +\code{"cuda:1"}) or a \code{torch_device}. A spec without an index +accepts any card; one with an index accepts only that card.} } \value{ -TRUE when every pair's live tensor is on the device's type; - FALSE on any mismatch or unreadable pair. Vacuously TRUE for an - empty staging set, which holds nothing to move. +TRUE when every pair's live tensor is on that device; FALSE + on any mismatch or unreadable pair. Vacuously TRUE for an empty + staging set, which holds nothing to move. } \description{ The check a caller makes before skipping an onload. It asks EVERY