From 9854f2022515e6ef191b672f2b2470b4f01d82ba Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Wed, 12 Jun 2024 17:12:18 +0000
Subject: [PATCH] fix: set sharded true if WORLD_SIZE is set

---
 server/text_generation_server/cli.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/server/text_generation_server/cli.py b/server/text_generation_server/cli.py
index 430323bc..1f5f3232 100644
--- a/server/text_generation_server/cli.py
+++ b/server/text_generation_server/cli.py
@@ -44,6 +44,9 @@ def serve(
     otlp_endpoint: Optional[str] = None,
     max_input_tokens: Optional[int] = None,
 ):
+    # derive sharded from environment variables if not provided
+    sharded = sharded or os.getenv("WORLD_SIZE", None) is not None
+
     if sharded:
         assert (
             os.getenv("RANK", None) is not None