From 2868accc407f92158a680cc1767a79526bf138c9 Mon Sep 17 00:00:00 2001
From: Travis Addair <travis@predibase.com>
Date: Fri, 1 Nov 2024 11:33:07 -0700
Subject: [PATCH] Docker test

---
 .github/workflows/build.yaml | 6 ++----
 router/src/batch.rs          | 1 -
 router/src/infer.rs          | 6 ------
 3 files changed, 2 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 13b9e96ca..3e7ac15a0 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - 'main'
+      - 'optimizations'
     tags:
       - 'v*'
 
@@ -69,10 +70,7 @@ jobs:
           images: |
             ghcr.io/predibase/lorax
           tags: |
-            type=semver,pattern={{version}}
-            type=semver,pattern={{major}}.{{minor}}
-            type=sha,prefix=,suffix=,format=short
-            type=raw,value=main,enable=${{ github.ref == 'refs/heads/main' }}
+            type=raw,value=optimizations-5,enable=${{ github.ref == 'refs/heads/optimizations' }}
 
       - name: Create a hash from tags
         env:
diff --git a/router/src/batch.rs b/router/src/batch.rs
index 2363dbdf6..39f8d32c7 100644
--- a/router/src/batch.rs
+++ b/router/src/batch.rs
@@ -222,7 +222,6 @@ impl BatchEntriesState {
         // TODO(travis): clone is not ideal, find a way to do this cleanly in place
         for r in self.batch_requests.clone().into_iter().rev() {
             let id = r.id;
-            tracing::info!("!!! drain::remove entry id={id:?}");
             let entry = self.batch_entries.remove(&id).unwrap();
             let adapter_index = r.adapter_index;
             let adapter = self.index_to_adapter.get_mut(&adapter_index).unwrap();
diff --git a/router/src/infer.rs b/router/src/infer.rs
index b3b608dc3..0a4c4e5cc 100644
--- a/router/src/infer.rs
+++ b/router/src/infer.rs
@@ -1323,7 +1323,6 @@ fn filter_send_generations(generations: Vec<Generation>, entries: &mut IntMap<u6
         let id = generation.request_id;
         // Get entry
         // We can `expect` here as the request id should always be in the entries
-        tracing::info!("!!! filter_send_generations id={id:?}");
         let entry = entries
             .get(&id)
             .expect("ID not found in entries. This is a bug.");
@@ -1338,7 +1337,6 @@ fn filter_send_generations(generations: Vec<Generation>, entries: &mut IntMap<u6
             metrics::increment_counter!("lorax_request_failure", "err" => "dropped");
         }).unwrap_or(true);
         if stopped {
-            tracing::info!("!!! filter_send_generations::remove entry id={id:?}");
             entries.remove(&id).expect("ID not found in entries. This is a bug.");
         }
     });
@@ -1351,7 +1349,6 @@ fn send_responses(
 ) -> Result<bool, Box<SendError<Result<InferStreamResponse, InferError>>>> {
     // Return directly if the channel is disconnected
     if entry.response_tx.is_closed() {
-        tracing::info!("!!! send_responses::disconnected");
         return Ok(true);
     }
 
@@ -1410,9 +1407,6 @@ fn send_responses(
 
         match (&generation.generated_text, iterator.peek()) {
             (Some(generated_text), None) => {
-                tracing::info!(
-                    "!!! send_responses::generation_ended id={id:?} generated_text={generated_text:?}"
-                );
                 // Generation has ended
                 stopped = true;
                 // Send message