From 68455353f5c158baf5576da6fcb670404ee0c391 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <olivier@huggingface.co>
Date: Wed, 15 Feb 2023 16:23:45 +0100
Subject: [PATCH] feat(launcher): add disable_custom_kernels arg (#67)

---
 launcher/src/main.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/launcher/src/main.rs b/launcher/src/main.rs
index ecb2c0b..8b0b55f 100644
--- a/launcher/src/main.rs
+++ b/launcher/src/main.rs
@@ -47,6 +47,8 @@ struct Args {
     #[clap(long, env)]
     weights_cache_override: Option<String>,
     #[clap(long, env)]
+    disable_custom_kernels: bool,
+    #[clap(long, env)]
     json_output: bool,
     #[clap(long, env)]
     otlp_endpoint: Option<String>,
@@ -79,6 +81,7 @@ fn main() -> ExitCode {
         master_port,
         huggingface_hub_cache,
         weights_cache_override,
+        disable_custom_kernels,
         json_output,
         otlp_endpoint,
     } = args;
@@ -242,6 +245,7 @@ fn main() -> ExitCode {
                 master_port,
                 huggingface_hub_cache,
                 weights_cache_override,
+                disable_custom_kernels,
                 otlp_endpoint,
                 status_sender,
                 shutdown,
@@ -406,6 +410,7 @@ fn shard_manager(
     master_port: usize,
     huggingface_hub_cache: Option<String>,
     weights_cache_override: Option<String>,
+    disable_custom_kernels: bool,
     otlp_endpoint: Option<String>,
     status_sender: mpsc::Sender<ShardStatus>,
     shutdown: Arc<Mutex<bool>>,
@@ -474,6 +479,11 @@ fn shard_manager(
         ));
     };
 
+    // If disable_custom_kernels is true, pass it to the shard as an env var
+    if disable_custom_kernels {
+        env.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
+    }
+
     // If the NCCL_SHM_DISABLE env var is set, pass it to the shard
     // needed when running NCCL inside a docker container and when you can't increase shm size
     if let Ok(nccl_shm_disalbe) = env::var("NCCL_SHM_DISABLE") {