1111# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212# See the License for the specific language governing permissions and
1313# limitations under the License.
14- from logging import Logger
1514import os
15+ import sys
16+ from logging import Logger
1617from threading import Event , Thread
1718from time import perf_counter , sleep
1819from typing import Optional
19- import sys
20+
2021
2122# Add the parent directory to Python path to import benchmarks_entrypoint
2223sys .path .insert (0 , os .path .dirname (os .path .dirname (os .path .abspath (__file__ ))))
23- from benchmarks_entrypoint import MetricsRecorder
24-
2524import gpustat
2625import psutil
2726import psycopg2
27+ from benchmarks_entrypoint import MetricsRecorder
28+
2829
2930# Optional heavy ML dependencies - only required when actually running the benchmark
3031try :
3132 import torch
33+
3234 from transformers import AutoModelForCausalLM , AutoTokenizer , GenerationConfig , StaticCache
35+
3336 TRANSFORMERS_AVAILABLE = True
3437except ImportError :
3538 TRANSFORMERS_AVAILABLE = False
@@ -63,19 +66,25 @@ def collect_metrics(benchmark_id, continue_metric_collection, metrics_recorder):
6366
6467
6568def run_benchmark (
66- logger : Logger , repository : str , branch : str , commit_id : str , commit_msg : str , metrics_recorder = None , num_tokens_to_generate = 100
69+ logger : Logger ,
70+ repository : str ,
71+ branch : str ,
72+ commit_id : str ,
73+ commit_msg : str ,
74+ metrics_recorder = None ,
75+ num_tokens_to_generate = 100 ,
6776):
6877 # Check if required ML dependencies are available
6978 if not TRANSFORMERS_AVAILABLE :
7079 logger .error ("Transformers and torch are required to run the LLaMA benchmark. Please install them with:" )
7180 logger .error ("pip install torch transformers" )
7281 logger .error ("Skipping LLaMA benchmark due to missing dependencies." )
7382 return
74-
83+
7584 continue_metric_collection = Event ()
7685 metrics_thread = None
7786 model_id = "meta-llama/Llama-2-7b-hf"
78-
87+
7988 # If no metrics_recorder is provided, create one for backward compatibility
8089 if metrics_recorder is None :
8190 try :
@@ -154,7 +163,7 @@ def sample(logits, temperature: float = 1.0, top_k: Optional[int] = None):
154163 # First eager forward pass
155164 logger .info ("running first eager forward pass" )
156165 start = perf_counter ()
157- outputs = model (** inputs )
166+ _ = model (** inputs )
158167 torch .cuda .synchronize ()
159168 end = perf_counter ()
160169 first_eager_fwd_pass_time = end - start
@@ -163,7 +172,7 @@ def sample(logits, temperature: float = 1.0, top_k: Optional[int] = None):
163172 # Second eager forward pass (should be faster)
164173 logger .info ("running second eager forward pass" )
165174 start = perf_counter ()
166- outputs = model (** inputs )
175+ _ = model (** inputs )
167176 torch .cuda .synchronize ()
168177 end = perf_counter ()
169178 second_eager_fwd_pass_time = end - start
@@ -339,7 +348,7 @@ def sample(logits, temperature: float = 1.0, top_k: Optional[int] = None):
339348 continue_metric_collection .set ()
340349 if metrics_thread is not None :
341350 metrics_thread .join ()
342-
351+
343352 # Only close the recorder if we created it locally
344353 if should_close_recorder :
345- metrics_recorder .close ()
354+ metrics_recorder .close ()
0 commit comments