change path to save file

tenstorrent · Nov 19, 2024 · 2583311 · 2583311
1 parent ceb3ce7
commit 2583311
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/tests/mock_vllm_model.py b/tests/mock_vllm_model.py
@@ -5,6 +5,8 @@
 import json
 from dataclasses import dataclass
 from typing import List
+from pathlib import Path
+
 
 import torch
 
@@ -342,7 +344,8 @@ def __init__(self, num_scheduler_steps, batch_size) -> None:
         self.num_scheduler_steps = num_scheduler_steps
         self.batch_size = batch_size
         timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
-        self.filepath = f"/home/user/tests/statistics_{timestamp}.jsonl"
+        cache_root = Path(os.getenv("CACHE_ROOT", "."))
+        self.filepath = cache_root / f"statistics_{timestamp}.jsonl"
         self.num_total_grouped_step = (
             0  # number of iterations of size num_scheduler_steps
         )

diff --git a/vllm-tt-metal-llama3-70b/src/logging_utils.py b/vllm-tt-metal-llama3-70b/src/logging_utils.py
@@ -1,6 +1,8 @@
 import os
 import datetime
 import json
+from pathlib import Path
+
 
 from vllm.engine.metrics_types import StatLoggerBase, Stats, SupportsMetricsInfo
 from vllm.engine.metrics import logger
@@ -25,7 +27,8 @@ def __init__(self, num_scheduler_steps, batch_size) -> None:
         self.num_scheduler_steps = num_scheduler_steps
         self.batch_size = batch_size
         timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
-        self.filepath = f"/home/user/tests/statistics_{timestamp}.jsonl"
+        cache_root = Path(os.getenv("CACHE_ROOT", "."))
+        self.filepath = cache_root / f"statistics_{timestamp}.jsonl"
         self.num_total_grouped_step = (
             0  # number of iterations of size num_scheduler_steps
         )