Investigate memory consumption on reading in CSV

[rok-cesnovar]

We have trimmed down the reading-in wall time quite a bit in the last few months, now we also need to investigate and profile the memory footprint of both read_cmdstan_csv and the $draws() function.

It seems that we are currently not near the optimal there, as can be seen in this report by @yizhang-yiz.

The baseline footprint is

fread_cmd <- paste0("grep -v '^#' --color=never ", output_file)
 draws <- data.table::fread(
          cmd = fread_cmd,
          data.table = FALSE
 )

[yizhang-yiz]

FWIW this is the output with verbose on:

> fread_cmd <- paste0("grep -v '^#' --color=never ", "*.csv")
> draws <- data.table::fread(cmd=fread_cmd, data.table=FALSE, verbose=TRUE)
  OpenMP version (_OPENMP)       201511
  omp_get_num_procs()            4
  R_DATATABLE_NUM_PROCS_PERCENT  unset (default 50)
  R_DATATABLE_NUM_THREADS        unset
  R_DATATABLE_THROTTLE           unset (default 1024)
  omp_get_thread_limit()         2147483647
  omp_get_max_threads()          4
  OMP_THREAD_LIMIT               unset
  OMP_NUM_THREADS                unset
  RestoreAfterFork               true
  data.table is using 2 threads with throttle==1024. See ?setDTthreads.
			Input contains no \n. Taking this to be a filename to open
[01] Check arguments
  Using 2 threads (omp_get_max_threads()=4, nth=2)
  NAstrings = [<<NA>>]
  None of the NAstrings look like numbers.
  show progress = 1
  0/1 column will be read as integer
[02] Opening the file
  Opening file /tmp/RtmpeZdCIY/file3bed6548c9ef
  File opened, size = 1.945GB (2088953363 bytes).
  Memory mapped ok
[03] Detect and skip BOM
[04] Arrange mmap to be \0 terminated
  \n has been found in the input and different lines can end with different line endings (e.g. mixed \n and \r\n in one file). This is common and ideal.
[05] Skipping initial rows if needed
  Positioned on line 1 starting: <<output.1.csv:lp__,accept_stat_>>
[06] Detect separator, quoting rule, and ncolumns
  Detecting sep automatically ...
  sep=','  with 100 lines of 68611 fields using quote rule 0
  Detected 68611 columns on line 1. This line is either column names or first data row. Line starts as: <<output.1.csv:lp__,accept_stat_>>
  Quote rule picked = 0
  fill=false and the most number of columns found is 68611
[07] Detect column types, good nrow estimate and whether first row is column names
  Number of sampling jump points = 10 because (2088953362 bytes from row 1 to eof) / (2 * 52480497 jump0size) == 19
  Type codes (jump 000)    : C7755577777777777775777775777775777775777775777777777777777777777777777777777777...7755557777  Quote rule 0
  Type codes (jump 001)    : C7755577777777777775777775777775777775777775777777777777777777777777777777777777...7755557777  Quote rule 0
  Type codes (jump 002)    : C7755577777777777775777775777775777775777775777777777777777777777777777777777777...7755557777  Quote rule 0
  Type codes (jump 003)    : C7755577777777777775777775777775777775777775777777777777777777777777777777777777...7755557777  Quote rule 0
  Type codes (jump 004)    : C7755577777777777775777775777775777775777775777777777777777777777777777777777777...7755557777  Quote rule 0
  Type codes (jump 005)    : CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC...CCCCCCCCCC  Quote rule 0
  Type codes (jump 010)    : CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC...CCCCCCCCCC  Quote rule 0
  'header' determined to be true because all columns are type string and a better guess is not possible
  =====
  Sampled 1050 rows (handled \n inside quoted fields) at 11 jump points
  Bytes from first data row on line 2 to the end of last row: 2088091035
  Line length: mean=521699.69 sd=-nan min=520746 max=862327
  Estimated number of rows: 2088091035 / 521699.69 = 4003
  Initial alloc = 4403 rows (4003 + 9%) using bytes/max(mean-2*sd,min) clamped between [1.1*estn, 2.0*estn]
  =====
[08] Assign column names
[09] Apply user overrides on column types
  After 0 type and 0 drop user overrides : CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC...CCCCCCCCCC
[10] Allocate memory for the datatable
  Allocating 68611 column slots (68611 - 0 dropped) with 4403 rows
[11] Read the data
  jumps=[0..4), chunk_size=522022758, total_size=2088091035
Read 4003 rows x 68611 columns from 1.945GB (2088953363 bytes) file in 01:58.742 wall clock time
[12] Finalizing the datatable
  Type counts:
     68611 : string    'C'
=============================
   0.000s (  0%) Memory map 1.945GB file
   2.509s (  2%) sep=',' ncol=68611 and header detection
   0.010s (  0%) Column type detection using 1050 sample rows
   4.313s (  4%) Allocation of 4403 rows x 68611 cols (2.252GB) of which 4003 ( 91%) rows used
 111.910s ( 94%) Reading 4 chunks (0 swept) of 497.840MB (each chunk 1000 rows) using 2 threads
   +    2.695s (  2%) Parse to row-major thread buffers (grown 0 times)
   +   83.922s ( 71%) Transpose
   +   25.293s ( 21%) Waiting
   0.000s (  0%) Rereading 0 columns due to out-of-sample type exceptions
 118.742s        Total
>
>
> object.size(draws)
16192859000 bytes

[rok-cesnovar]

Thanks! Can you tell me if all of output values are doubles or are there any integers as well? If there are both types what is the ratio between them?

[yizhang-yiz]

All output values are doubles.

[rok-cesnovar]

I am a bit baffled about why the 4x in size your case.

library(cmdstanr)
model_code <- '
data {
  int N;
}
generated quantities {
  real a[N];
  for(i in 1:N) {
     a[i] = 0.1234567891;
  }
}
'
stan_model <- write_stan_file(code = model_code)

m <- cmdstan_model("./generate_big_file.stan")

for (i in seq(1,10)) {
  utils::capture.output(
    f <- m$sample(data=list(N = 2000, val = 0),chains =1, fixed_param = TRUE, iter_sampling = 4000, sig_figs = i)
  )
  d <- f$draws()  
  print(as.numeric(object.size(d))/file.size(f$output_files()))
}

I get the following:

[1] 2.010154
[1] 1.608425
[1] 1.340522
[1] 1.149122
[1] 1.005549
[1] 0.8938678
[1] 0.8045146
[1] 0.7314019
[1] 0.6704709
[1] 0.6189111

Meaning that if all outputs are 0.1 the draws object is twice the size of the CSV, which is expected as each sample takes 4 bytes: an ascii character each for "0.1,". The larger the precision, the larger the CSV file while the draws object size remains the same. Which all makes sense.

Not sure what is going on.

[yizhang-yiz]

It baffles me too.

This dummy model gives approx same size of the CSV & number of parameters.

data {
 int N;
}

parameters {
 vector[N] x;
}

model {
 x ~ normal(0, 1);
}

fit <- mod$sample(data=list(N=65000), fixed_param=TRUE, iter_sampling=1000)
fit$save_output_files()

But the draws size is nowhere near the draws.original read from original CSVs.

> fread_cmd <- paste0("grep -v '^#' --color=never ", "*.csv")
>  draws <- data.table::fread(
+           cmd = fread_cmd,
+           data.table = FALSE
+  )
> object.size(draws)
2097069616 bytes
> fread_cmd <- paste0("grep -v '^#' --color=never ", "../output.[1-4].csv")
> draws.original <- data.table::fread(
+           cmd = fread_cmd,
+           data.table = FALSE
+  )
> object.size(draws.original)
16192859000 bytes

[rok-cesnovar]

Thanks for the files @yizhang-yiz

There are two separate things here.

1. fread should be used with 1 file at a time

   We do read 1 file at a time in cmdstanr but in our debugging we did not.

The following

fread_cmd <- paste0("grep -v '^#' --color=never ", "output.1.csv output.2.csv")
draws <- data.table::fread(
             cmd = fread_cmd,
             data.table = FALSE
         )

produces a 8GB object out of 2 0.5GB files. The first column is all weird with a bunch of "output.1.csv:"

If I run

fread_cmd <- paste0("grep -v '^#' --color=never ", "output.1.csv")
draws_1 <- data.table::fread(
  cmd = fread_cmd,
  data.table = FALSE
)
fread_cmd <- paste0("grep -v '^#' --color=never ", "output.2.csv")
draws_2 <- data.table::fread(
  cmd = fread_cmd,
  data.table = FALSE
)

and then rbind the results its 1GB. So this is a fread thing (or maybe thats how grep works with multiple files idk), may be a bug or not, not sure. But this isnt whats happening in cmdstanr.

2. the returned draws object is the expected size but we use too much memory internally

> fit <- cmdstanr::as_cmdstan_fit(c("output.1.csv", "output.2.csv"))
> draws_fit <- fit$draws()
> object.size(draws_fit)
1102868328 bytes # this is ~1GB

However, observing system monitor I saw spikes of RAM usage up to 6GB, so seems we do some copies somewhere.

The issue is in read_cmdstan_csv.

[yizhang-yiz]

I see. Thanks a lot. I think the 2nd item is what caused issue in my original importing process, as reading all 4 chains

fit <- cmdstanr::as_cmdstan_fit(c("output.1.csv", "output.2.csv", ""output.3.csv", "output.4.csv"))

shows RAM usage to ~16G on system monitor. But apparently object.size reporting similar figure during debug is just a coincidence.

[rok-cesnovar]

I have stared at this for a bit and am getting closer. We can do a few things in cmdstanr to be a tad more efficient but the biggest culprit is posterior::bind_draws.

This code replicates what we do in read_cmdstan_csv:

data_file <- c(file.path("output.1.csv"), file.path("output.2.csv"), file.path("output.3.csv"), file.path("output.4.csv"))

draws <- list()
i <- 1
for (f in data_file) {
  fread_cmd <- paste0("grep -v '^#' --color=never ", f)
  draws[[i]] <- data.table::fread(
    cmd = fread_cmd,
    data.table = FALSE
  )
  i <- i + 1
}
draws <- lapply(draws, posterior::as_draws_array)
draws[["along"]] <- "chain"
draws <- do.call(posterior::bind_draws, draws)
gc()

Below is the RAM usage. The circled area is the call to bind_draws.

[rok-cesnovar]

I have fixes for this, will open a PR tomorrow/wednesday.

[bachlaw]

Not sure if this is the same issue, but I have noted on a Linux installation that cmdstanr, at least when used with brms, seems to be consuming (and not giving back) a lot of memory, ~ 2 GB each time, to the point where I am eventually informed no more models can run because forking is (no longer) possible.

Stan usually is very good about not absorbing memory. I had assumed this was just the opportunity cost of reduce_sum threading (here: 1 core, 7 threads), but I wonder if what you are describing here is the real mechanism.

[rok-cesnovar]

Unfortunately this is not related to your issue. The thing in this issue causes RAM usage after the sampling finishes and you use fit$draw() or read_cmdstan_csv() or any other cmdstanr function that reads in samples. brms uses rstan::read_stan_csv() after cmdstanr finishes to read the samples.