11import argparse
22import time
3- from typing import Any , List , Optional
43
54from rich .live import Live
6- from rich .table import Table
75
86from dstack ._internal .cli .commands import APIBaseCommand
97from dstack ._internal .cli .services .completion import RunNameCompleter
108from dstack ._internal .cli .utils .common import (
11- LIVE_TABLE_PROVISION_INTERVAL_SECS ,
129 LIVE_TABLE_REFRESH_RATE_PER_SEC ,
13- add_row_from_dict ,
1410 console ,
1511)
12+ from dstack ._internal .cli .utils .metrics import (
13+ MAX_SAMPLES ,
14+ WATCH_INTERVAL_SECONDS ,
15+ get_metrics_table ,
16+ )
1617from dstack ._internal .core .errors import CLIError
17- from dstack ._internal .core .models .instances import Resources
1818from dstack ._internal .core .models .metrics import JobMetrics
19+ from dstack ._internal .core .models .runs import Job
1920from dstack .api ._public import Client
2021from dstack .api ._public .runs import Run
2122
@@ -33,121 +34,61 @@ def _register(self):
3334 help = "Watch run metrics in realtime" ,
3435 action = "store_true" ,
3536 )
37+ self ._parser .add_argument (
38+ "--replica" ,
39+ help = "The replica number. Defaults to 0." ,
40+ type = int ,
41+ default = 0 ,
42+ )
43+ self ._parser .add_argument (
44+ "--job" ,
45+ help = "The job number inside the replica. Defaults to 0." ,
46+ type = int ,
47+ default = 0 ,
48+ )
3649
3750 def _command (self , args : argparse .Namespace ):
3851 super ()._command (args )
39- run = self .api .runs .get (run_name = args .run_name )
40- if run is None :
41- raise CLIError (f"Run { args .run_name } )
42- metrics = _get_run_jobs_metrics (api = self .api , run = run )
52+ job , metrics = self ._fetch (args )
4353
4454 if not args .watch :
45- console .print (_get_metrics_table ( run , metrics ))
55+ console .print (get_metrics_table ( job , metrics ))
4656 return
4757
4858 try :
4959 with Live (console = console , refresh_per_second = LIVE_TABLE_REFRESH_RATE_PER_SEC ) as live :
5060 while True :
51- live .update (_get_metrics_table (run , metrics ))
52- time .sleep (LIVE_TABLE_PROVISION_INTERVAL_SECS )
53- run = self .api .runs .get (run_name = args .run_name )
54- if run is None :
55- raise CLIError (f"Run { args .run_name } )
56- metrics = _get_run_jobs_metrics (api = self .api , run = run )
61+ live .update (get_metrics_table (job , metrics ))
62+ time .sleep (WATCH_INTERVAL_SECONDS )
63+ job , metrics = self ._fetch (args )
5764 except KeyboardInterrupt :
5865 pass
5966
67+ def _fetch (self , args : argparse .Namespace ) -> tuple [Job , JobMetrics ]:
68+ run = self .api .runs .get (run_name = args .run_name )
69+ if run is None :
70+ raise CLIError (f"Run { args .run_name } )
71+ job = _get_job (run , args .replica , args .job )
72+ return job , _get_job_metrics (self .api , run , job )
6073
61- def _get_run_jobs_metrics (api : Client , run : Run ) -> List [JobMetrics ]:
62- metrics = []
63- for job in run ._run .jobs :
64- job_metrics = api .client .metrics .get_job_metrics (
65- project_name = api .project ,
66- run_name = run .name ,
67- replica_num = job .job_spec .replica_num ,
68- job_num = job .job_spec .job_num ,
69- )
70- metrics .append (job_metrics )
71- return metrics
72-
73-
74- def _get_metrics_table (run : Run , metrics : List [JobMetrics ]) -> Table :
75- table = Table (box = None )
76- table .add_column ("NAME" , style = "bold" , no_wrap = True )
77- table .add_column ("STATUS" )
78- table .add_column ("CPU" )
79- table .add_column ("MEMORY" )
80- table .add_column ("GPU" )
81-
82- run_row = {"NAME" : run .name , "STATUS" : run .status .value }
83- if len (run ._run .jobs ) != 1 :
84- add_row_from_dict (table , run_row )
85-
86- for job , job_metrics in zip (run ._run .jobs , metrics ):
87- jrd = job .job_submissions [- 1 ].job_runtime_data
88- jpd = job .job_submissions [- 1 ].job_provisioning_data
89- resources : Optional [Resources ] = None
90- if jrd is not None and jrd .offer is not None :
91- resources = jrd .offer .instance .resources
92- elif jpd is not None :
93- resources = jpd .instance_type .resources
94- cpu_usage = _get_metric_value (job_metrics , "cpu_usage_percent" )
95- if cpu_usage is not None :
96- if resources is not None :
97- cpu_usage = cpu_usage / resources .cpus
98- cpu_usage = f"{ cpu_usage :.0f}
99- memory_usage = _get_metric_value (job_metrics , "memory_working_set_bytes" )
100- if memory_usage is not None :
101- memory_usage = _format_memory (memory_usage , 2 )
102- if resources is not None :
103- memory_usage += f"/{ _format_memory (resources .memory_mib * 1024 * 1024 , 2 )}
104- gpu_metrics = ""
105- gpus_detected_num = _get_metric_value (job_metrics , "gpus_detected_num" )
106- if gpus_detected_num is not None :
107- for i in range (gpus_detected_num ):
108- gpu_memory_usage = _get_metric_value (job_metrics , f"gpu_memory_usage_bytes_gpu{ i } )
109- gpu_util_percent = _get_metric_value (job_metrics , f"gpu_util_percent_gpu{ i } )
110- if gpu_memory_usage is not None :
111- if i != 0 :
112- gpu_metrics += "\n "
113- gpu_metrics += f"gpu={ i } { _format_memory (gpu_memory_usage , 2 )}
114- if resources is not None :
115- gpu_metrics += (
116- f"/{ _format_memory (resources .gpus [i ].memory_mib * 1024 * 1024 , 2 )}
117- )
118- gpu_metrics += f" util={ gpu_util_percent }
119-
120- job_row = {
121- "NAME" : f" replica={ job .job_spec .replica_num } { job .job_spec .job_num } ,
122- "STATUS" : job .job_submissions [- 1 ].status .value ,
123- "CPU" : cpu_usage or "-" ,
124- "MEMORY" : memory_usage or "-" ,
125- "GPU" : gpu_metrics or "-" ,
126- }
127- if len (run ._run .jobs ) == 1 :
128- job_row .update (run_row )
129- add_row_from_dict (table , job_row )
130-
131- return table
132-
133-
134- def _get_metric_value (job_metrics : JobMetrics , name : str ) -> Optional [Any ]:
135- for metric in job_metrics .metrics :
136- if metric .name == name :
137- return metric .values [- 1 ]
138- return None
139-
140-
141- def _format_memory (memory_bytes : int , decimal_places : int ) -> str :
142- """See test_format_memory in tests/_internal/cli/commands/test_metrics.py for examples."""
143- memory_mb = memory_bytes / 1024 / 1024
144- if memory_mb >= 1024 :
145- value = memory_mb / 1024
146- unit = "GB"
147- else :
148- value = memory_mb
149- unit = "MB"
15074
151- if decimal_places == 0 :
152- return f"{ round (value )} { unit }
153- return f"{ value :.{decimal_places }f} .rstrip ("0" ).rstrip ("." ) + unit
75+ def _get_job (run : Run , replica_num : int , job_num : int ) -> Job :
76+ for job in run ._run .jobs :
77+ if job .job_spec .replica_num == replica_num and job .job_spec .job_num == job_num :
78+ return job
79+ raise CLIError (
80+ f"Run { run .name } { replica_num } { job_num }
81+ " Use --replica and --job to select one."
82+ )
83+
84+
85+ def _get_job_metrics (api : Client , run : Run , job : Job ) -> JobMetrics :
86+ """`limit` must be sent explicitly: the endpoint declares it `limit: int = 1`, not
87+ Optional, so omitting it caps the response at one sample."""
88+ return api .client .metrics .get_job_metrics (
89+ project_name = api .project ,
90+ run_name = run .name ,
91+ replica_num = job .job_spec .replica_num ,
92+ job_num = job .job_spec .job_num ,
93+ limit = MAX_SAMPLES ,
94+ )
0 commit comments