Simplify to use task queue

Dockerfile

@@ -22,13 +22,9 @@ RUN installGithub.r \
     mlverse/tok
     
 RUN installGithub.r \
-    mlverse/minhub
-
-RUN installGithub.r \
-    mlverse/hfhub 
-
-RUN installGithub.r \
-    mlverse/minhub
+    mlverse/minhub \
+    mlverse/hfhub \
+    mlverse/callq
 
 # see: https://huggingface.co/docs/hub/spaces-sdks-docker#permissions
 RUN useradd -m -u 1000 user

app.R

@@ -4,16 +4,10 @@ library(minhub)
 library(magrittr)
 source("model-session.R")
 
-repo <- "stabilityai/stablelm-tuned-alpha-3b"
+repo <- "EleutherAI/pythia-70m"
 repo <- Sys.getenv("MODEL_REPO", unset = repo)
 sess <- model_session$new()
 
-poll_process <- function() {
-  sess$poll_process(1)
-  later::later(func = poll_process, delay = 0.5)
-}
-poll_process()
-
 max_n_tokens <- 100
 system_prompt = "<|SYSTEM|># StableLM Tuned (Alpha version)
 - StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
@@ -124,7 +118,6 @@ server <- function(input, output, session) {
     
     # this runs for the cases where sess$is_loaded was NULL
     # ie there was no model currently loading.
-    cat("Loading model:",sess$sess$poll_process(), "\n")
     m <- model_loaded() %>% 
       promises::then(onFulfilled = function(x) {
         cat("Model has been loaded!", "\n")

model-session.R

@@ -1,28 +1,20 @@
-source("promise-session.R")
 
-# A wrapper a around the promise session that controls model loading and
-# querying given a prompt
 model_session <- R6::R6Class(
   lock_objects = FALSE,
   public = list(
     initialize = function() {
-      self$sess <- NULL
+      self$task_q <- NULL
       self$temperature <- 1
       self$top_k <- 50
       self$is_loaded <- NULL
     },
-    poll_process = function(timeout = 1) {
-      if (!is.null(self$sess)) {
-        self$sess$poll_process(timeout)
-      }
-    },
     load_model = function(repo) {
       if (!is.null(self$sess)) {
         cat("Model is already loaded.", "\n")
-        return(self$sess$call(function() "done"))
+        return(self$task_q$push(function() "done"))
       }
-      self$sess <- promise_session$new()
-      self$sess$call(args = list(repo = repo), function(repo) {
+      self$task_q <- callq::task_q$new(num_workers = 1)
+      self$task_q$push(args = list(repo = repo), function(repo) {
         library(torch)
         library(zeallot)
         library(minhub)
@@ -34,16 +26,16 @@ model_session <- R6::R6Class(
       })
     },
     generate = function(prompt) {
-      if (is.null(self$sess)) {
+      if (is.null(self$task_q)) {
         cat("Model is not loaded, error.", "\n")
-        return(self$sess$call(function() stop("Model is not loaded")))
+        return(self$task_q$push(function() stop("Model is not loaded")))
       }
       args <- list(
         prompt = prompt, 
         temperature = self$temperature,
         top_k = self$top_k
       )
-      self$sess$call(args = args, function(prompt, temperature, top_k) {
+      self$task_q$push(args = args, function(prompt, temperature, top_k) {
         idx <- torch_tensor(tok$encode(prompt)$ids)$view(c(1, -1))
         with_no_grad({
           logits <- model(idx + 1L)

promise-session.R

@@ -1,70 +0,0 @@
-
-# Small utility class that wraps a `callr::r_session` to return promises when
-# executing `sess$call()`.
-# Only one promise is resolve per time in fifo way.
-promise_session <- R6::R6Class(
-  lock_objects = FALSE,
-  public = list(
-    initialize = function() {
-      self$sess <- callr::r_session$new()
-      self$is_running <- FALSE
-    },
-    call = function(func, args = list()) {
-      self$poll_process()
-      promises::promise(function(resolve, reject) {
-        self$push_task(func, args, resolve, reject)
-        later::later(self$poll_process, 1)
-      })
-    },
-    push_task = function(func, args, resolve, reject) {
-      self$tasks[[length(self$tasks) + 1]] <- list(
-        func = func, 
-        args = args, 
-        resolve = resolve, 
-        reject = reject
-      )
-      cat("task pushed, now we have ", length(self$tasks), " on queue\n")
-      self$run_task()
-      invisible(NULL)
-    },
-    run_task = function() {
-      if (self$is_running) return(NULL)
-      if (length(self$tasks) == 0) return(NULL)
-      
-      self$is_running <- TRUE
-      task <- self$tasks[[1]]
-      self$sess$call(task$func, args = task$args)
-    },
-    resolve_task = function() {
-      cat("Resolving task! ")
-      out <- self$sess$read()
-      if (!is.null(out$error)) {
-        self$tasks[[1]]$reject(out$error)
-      } else {
-        self$tasks[[1]]$resolve(out$result)
-      }
-      
-      self$tasks <- self$tasks[-1]
-      cat("now we have ", length(self$tasks), "on queue\n")
-      
-      self$is_running <- FALSE
-      
-      self$run_task()
-    },
-    poll_process = function(timeout = 1) {
-      if (!self$is_running) return("ready")
-      poll_state <- self$sess$poll_process(timeout)
-      if (poll_state == "ready") {
-        self$resolve_task()
-      }
-      poll_state
-    }
-  )
-)
-
-# sess <- promise_session$new()
-# f <- sess$call(function(a) {
-#   10 + 1
-# }, list(1))
-# sess$poll_process()
-