refactor(sglang): replace magic strings with constants and cache block shape

- Extract mode literals ("local"/"distributed") into module-level constants
  MODE_LOCAL, MODE_DISTRIBUTED, _VALID_MODES to prevent typo-induced bugs
- Extract error operation labels ("get"/"set"/"exists") into _OP_GET,
  _OP_SET, _OP_EXISTS constants for consistent Prometheus label usage
- Cache block shape tuple as self._block_shape at init time instead of
  recomputing kv_dim and constructing the shape on every _get_block_shaped()
  call (hot path in batch_get_v1 per-block loop)
- Remove dead field _started (set but never read)
- Update tests to import and use the new constants

Signed-off-by: staryxchen <staryxchen@tencent.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: staryxchen

flexkv/integration/sglang/hicache_storage_adapter.py

@@ -46,6 +46,19 @@
 except ImportError:
     from sglang.srt.observability.metrics_collector import StorageMetrics
 
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+MODE_LOCAL = "local"
+MODE_DISTRIBUTED = "distributed"
+_VALID_MODES = (MODE_LOCAL, MODE_DISTRIBUTED)
+
+# Error operation labels (used in metrics recording)
+_OP_GET = "get"
+_OP_SET = "set"
+_OP_EXISTS = "exists"
+
 
 # ---------------------------------------------------------------------------
 # Helper: extract token_ids from extra_info
@@ -162,17 +175,17 @@ def __init__(
         self._cache_config = CacheConfig(**cache_kwargs)
 
         # Extract distributed mode configuration
-        self._mode: str = extra.get("mode", "local")
-        if self._mode not in ("local", "distributed"):
-            raise ValueError(f"Invalid mode: {self._mode}. Must be 'local' or 'distributed'")
+        self._mode: str = extra.get("mode", MODE_LOCAL)
+        if self._mode not in _VALID_MODES:
+            raise ValueError(f"Invalid mode: {self._mode}. Must be one of {_VALID_MODES}")
 
         self._redis_host: str = extra.get("redis_host", "127.0.0.1")
         self._redis_port: int = extra.get("redis_port", 6379)
         self._redis_password: Optional[str] = extra.get("redis_password", None)
 
         self._prefetch_timeout: float = float(extra.get("prefetch_timeout", 5.0))
 
-        if self._mode == "distributed" and not self._redis_host:
+        if self._mode == MODE_DISTRIBUTED and not self._redis_host:
             raise ValueError("redis_host is required when mode='distributed'")
 
         self._should_backup: bool = (
@@ -184,9 +197,9 @@ def __init__(
         self._kv_manager = None
         self._cpu_cache_tensor = None  # direct access to CPU cache (thread mode)
         self._elements_per_block: int = 0
+        self._block_shape: tuple = ()  # set after KVManager init
         self._page_size: int = self._cache_config.tokens_per_block
         self._mem_pool_host = mem_pool_host
-        self._started = False
         self._bytes_per_page: int = 0
         self._gb_per_page: float = 0.0
 
@@ -231,8 +244,7 @@ def _init_kv_manager(self):
         }
 
         # Branch on mode: local vs distributed
-        if self._mode == "distributed":
-            # Distributed mode: enable cross-node KV Cache sharing
+        if self._mode == MODE_DISTRIBUTED:
             cache_config_kwargs.update({
                 "enable_remote": True,
                 "enable_kv_sharing": True,
@@ -273,13 +285,18 @@ def _init_kv_manager(self):
         # Get direct access to CPU cache tensor (thread mode only)
         self._cpu_cache_tensor = self._kv_manager.get_cpu_cache_tensor()
 
-        # Compute elements per block for indexing
+        # Compute elements per block and cache the block shape for indexing
         kv_dim = 1 if self._model_config.use_mla else 2
         self._elements_per_block = (
             self._model_config.num_layers * kv_dim *
             self._page_size * self._model_config.num_kv_heads *
             self._model_config.head_size
         )
+        self._block_shape = (
+            self._model_config.num_layers, kv_dim,
+            self._page_size, self._model_config.num_kv_heads,
+            self._model_config.head_size
+        )
 
         # Compute bytes_per_page for bandwidth reporting
         dtype = getattr(self._model_config, 'dtype', None) or torch.float16
@@ -346,12 +363,7 @@ def _get_block_view(self, block_id: int) -> torch.Tensor:
 
     def _get_block_shaped(self, block_id: int) -> torch.Tensor:
         """Get a CPU cache block reshaped to BLOCKFIRST: [L, kv_dim, T, H, D]."""
-        kv_dim = 1 if self._model_config.use_mla else 2
-        return self._get_block_view(block_id).view(
-            self._model_config.num_layers, kv_dim,
-            self._page_size, self._model_config.num_kv_heads,
-            self._model_config.head_size
-        )
+        return self._get_block_view(block_id).view(self._block_shape)
 
     def _fetch_remote_blocks(self, token_ids: np.ndarray) -> bool:
         """Fetch remote blocks into local CPU cache via prefetch_async.
@@ -411,7 +423,6 @@ def register_mem_pool_host(self, mem_pool_host: Any) -> None:
             self._page_size = sglang_page_size
 
         self._init_kv_manager()
-        self._started = True
 
         # KVManager has started — global collector is now initialized
         try:
@@ -499,7 +510,7 @@ def batch_exists(
                                     tokens_per_block=page_size)
 
             # In distributed mode, query both local and remote trees
-            if (self._mode == "distributed"
+            if (self._mode == MODE_DISTRIBUTED
                     and hasattr(cache_engine.cpu_cache_engine, 'match_all')):
                 match_result = cache_engine.cpu_cache_engine.match_all(seq_meta)
             else:
@@ -510,7 +521,7 @@ def batch_exists(
         except Exception:
             logger.exception("batch_exists failed")
             if self._metrics:
-                self._metrics.record_sglang_error("exists")
+                self._metrics.record_sglang_error(_OP_EXISTS)
             return 0
 
     def batch_get_v1(
@@ -550,7 +561,7 @@ def batch_get_v1(
                                     tokens_per_block=page_size)
 
             # In distributed mode, query both local and remote trees
-            if (self._mode == "distributed"
+            if (self._mode == MODE_DISTRIBUTED
                     and hasattr(cache_engine.cpu_cache_engine, 'match_all')):
                 match_result = cache_engine.cpu_cache_engine.match_all(seq_meta)
 
@@ -610,7 +621,7 @@ def batch_get_v1(
         except Exception:
             logger.exception("batch_get_v1 failed")
             if self._metrics:
-                self._metrics.record_sglang_error("get")
+                self._metrics.record_sglang_error(_OP_GET)
             return [False] * len(keys)
 
     def batch_set_v1(
@@ -694,7 +705,7 @@ def batch_set_v1(
         except Exception:
             logger.exception("batch_set_v1 failed")
             if self._metrics:
-                self._metrics.record_sglang_error("set")
+                self._metrics.record_sglang_error(_OP_SET)
             return [False] * len(keys)
 
     # Legacy abstract methods

flexkv/integration/sglang/test_hicache_storage_adapter.py

@@ -49,7 +49,7 @@
 from unittest.mock import MagicMock
 
 from flexkv.integration.sglang.hicache_storage_adapter import (
-    FlexKVHiCacheStorage, _get_token_ids,
+    FlexKVHiCacheStorage, _get_token_ids, MODE_LOCAL, MODE_DISTRIBUTED,
 )
 from sglang.srt.mem_cache.hicache_storage import (
     HiCacheStorageConfig, HiCacheStorageExtraInfo,
@@ -230,7 +230,7 @@ def test_no_token_ids_degradation():
 def test_default_local_mode():
     """Default mode is 'local' when not specified."""
     backend = FlexKVHiCacheStorage(_make_config())
-    assert backend._mode == "local", f"Expected 'local', got '{backend._mode}'"
+    assert backend._mode == MODE_LOCAL, f"Expected '{MODE_LOCAL}', got '{backend._mode}'"
 
 
 def test_explicit_local_mode():
@@ -247,8 +247,7 @@ def test_explicit_local_mode():
         }
     )
     backend = FlexKVHiCacheStorage(config)
-    assert backend._mode == "local"
-
+    assert backend._mode == MODE_LOCAL
 
 def test_distributed_mode_config():
     """Distributed mode stores redis config correctly."""
@@ -267,7 +266,7 @@ def test_distributed_mode_config():
         }
     )
     backend = FlexKVHiCacheStorage(config)
-    assert backend._mode == "distributed"
+    assert backend._mode == MODE_DISTRIBUTED
     assert backend._redis_host == "redis.example.com"
     assert backend._redis_port == 6380
     assert backend._redis_password == "test_password"