[*] Typo fixes

Author: SamuelMarks

benchmarks/mlperf/backend.py

@@ -327,7 +327,7 @@ def flush_queries(self):
         self.accuracy_log.write(json.dumps(pred_outputs))
         self.accuracy_log.flush()
         self.accuracy_log.close()
-        log.info("Dumpped prediction outputs to accuracy log... ")
+        log.info("Dumped prediction outputs to accuracy log... ")
 
   def __del__(self):
     print("Finished destroying SUT.")

benchmarks/tests/test_benchmark_serving.py

@@ -44,7 +44,7 @@ async def test_benchmark(self):
     disable_tqdm = True
 
     async def mocked_decode_response():
-      """Mocks decode reponse as an async generator."""
+      """Mocks decode response as an async generator."""
       responses = [
           jetstream_pb2.DecodeResponse(
               stream_content=jetstream_pb2.DecodeResponse.StreamContent(

docs/observability-prometheus-metrics-in-jetstream-server.md

@@ -1,6 +1,6 @@
 # Observability in JetStream Server
 
-In JetStream Server, we use [Prometheus](https://prometheus.io/docs/introduction/overview/) to collect key metrics within JetStream orchestrator and engines. We implemented a [Prometheus client server](https://prometheus.github.io/client_python/exporting/http/) in JetStream `server_lib.py` and use `MetricsServerConfig` (by passing `prometheus_port` in server entrypoint) to gaurd the metrics observability feature.
+In JetStream Server, we use [Prometheus](https://prometheus.io/docs/introduction/overview/) to collect key metrics within JetStream orchestrator and engines. We implemented a [Prometheus client server](https://prometheus.github.io/client_python/exporting/http/) in JetStream `server_lib.py` and use `MetricsServerConfig` (by passing `prometheus_port` in server entrypoint) to guard the metrics observability feature.
 
 ## Enable Prometheus server to observe Jetstream metrics
 

experimental/jax/inference/entrypoint/mini_offline_benchmarking.py

@@ -64,12 +64,12 @@ def benchmark():
   num_input_tokens = sum(map(lambda r: len(r.input_tokens), res_list))
   num_output_tokens = sum(map(lambda r: len(r.generated_tokens), res_list))
 
-  print("Benchmarking result: ")
-  print("  Total requests:", len(dataset))
-  print("  Total input tokens:", num_input_tokens)
-  print("  Total output tokens:", num_output_tokens)
-  print(f"  Input token thruput: {num_input_tokens/duration: .2f} tokens/sec")
-  print(f"  Output token thruput: {num_output_tokens/duration: .2f} tokens/sec")
+  print("Benchmarking result:")
+  print("  Total requests:          ", len(dataset))
+  print("  Total input tokens:      ", num_input_tokens)
+  print("  Total output tokens:     ", num_output_tokens)
+  print(f"  Input token throughput:  {num_input_tokens/duration: .2f} tokens/sec")
+  print(f"  Output token throughput: {num_output_tokens/duration: .2f} tokens/sec")
 
 
 if __name__ == "__main__":

experimental/jax/inference/nn/linear.py

@@ -7,7 +7,7 @@
 
      https://www.apache.org/licenses/LICENSE-2.0
 
-Unless reuired by applicable law or agreed to in writing, software
+Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and

experimental/jax/inference/parallel/operations.py

@@ -22,7 +22,7 @@
 
 
 def reduce_scatter(operand, scatter_dimension, axis_names):
-  """reduce-scatter sum operation via ppermute."""
+  """reduce-scatter sum operation via permute."""
   idx = get_partition_index(axis_names=axis_names)
   num_partitions = get_num_partitions(axis_names=axis_names)
   chunk_size = operand.shape[scatter_dimension] // num_partitions

experimental/jax/inference/parallel/util.py

@@ -27,6 +27,6 @@ def pspec(a):
     elif isinstance(a, int) or isinstance(a, float):
       return P()
     else:
-      raise ValueError(f"unknown parition spec for {a}")
+      raise ValueError(f"unknown partition spec for {a}")
 
   return jax.tree_util.tree_map(pspec, sharded_pytree)

experimental/jax/inference/runtime/batch_scheduler.py

@@ -110,20 +110,20 @@ def schedule(
         cur_prompt_chunk_len = (
             total_len - next_prefill_req.chunk_idx * next_prefill_req.chunk_size
         )
-      alloced_pages = self.kv_manager.alloc_prefill_hbm_pages(
+      allocated_pages = self.kv_manager.alloc_prefill_hbm_pages(
           cur_prompt_chunk_len
       )
-      if len(alloced_pages) == 0:
+      if len(allocated_pages) == 0:
         # TODO: introduce priority for the request and better
         # eviction algorithm.
         raise NotImplementedError("Eviction is not supported yet")
       else:
         start_idx = (
             next_prefill_req.chunk_idx * next_prefill_req.chunk_size
         ) // self.kv_manager.page_size
-        for i, page in enumerate(alloced_pages):
+        for i, page in enumerate(allocated_pages):
           next_prefill_req.page_indices[start_idx + i] = page
-        prefill_pages_update = PrefillPagesUpdate(alloced_pages)
+        prefill_pages_update = PrefillPagesUpdate(allocated_pages)
 
     # Schedule new generate reqs and allocate memory for all reqs.
     with generate_state.map_mutex:
@@ -150,12 +150,12 @@ def schedule(
           next_generate_reqs.append(gr)
 
         # Check and alloc memory for generate.
-        alloced_pages = self.kv_manager.alloc_hbm_pages(
+        allocated_pages = self.kv_manager.alloc_hbm_pages(
             len(generate_state.active_slot_req_map)
         )
         if (
             len(generate_state.active_slot_req_map) != 0
-            and len(alloced_pages) == 0
+            and len(allocated_pages) == 0
         ):
           raise NotImplementedError(
               "Eviction isn't supported yet, please set a lower value for batch_size"
@@ -169,17 +169,17 @@ def schedule(
           if idx >= len(req.page_indices):
             continue
 
-          req.page_indices[idx] = alloced_pages[page_to_use]
+          req.page_indices[idx] = allocated_pages[page_to_use]
           generate_state_page_updates.append(
               GenerateStatePageUpdate(
                   slot=slot,
                   page_idx=idx,
-                  mapped_idx=alloced_pages[page_to_use],
+                  mapped_idx=allocated_pages[page_to_use],
               )
           )
           page_to_use += 1
 
-        self.kv_manager.free_hbm_pages(alloced_pages[page_to_use:])
+        self.kv_manager.free_hbm_pages(allocated_pages[page_to_use:])
 
         if len(generate_state.active_slot_req_map) == 0:
           schedule_generate = False

experimental/jax/inference/runtime/engine.py

@@ -223,7 +223,7 @@ def __init__(
     )
     print(" preprocess,", end="")
     self._preprocess_queue: queue.Queue[Request] = queue.Queue()
-    # TODO: Seperate the running loop with the static inference model.
+    # TODO: Separate the running loop with the static inference model.
     self._preprocess_thread = threading.Thread(
         name="preprocess", target=self._preprocess
     )

experimental/jetstream-maxtext-stable-stack/build.sh

@@ -37,4 +37,4 @@ docker build --no-cache \
     -t ${LOCAL_IMAGE_TAG} \
     -f ./Dockerfile .
 
-echo "********* Sucessfully built Stable Stack Image with tag $LOCAL_IMAGE_TAG *********"
+echo "********* Successfully built Stable Stack Image with tag $LOCAL_IMAGE_TAG *********"