wehub-resource-sync
202 lines
7.2 KiB
Python
202 lines
7.2 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# 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 limitations
|
|
# under the License.
|
|
"""Configurations for measurements in the runner"""
|
|
|
|
import os
|
|
from threading import Thread
|
|
from typing import NamedTuple, Optional
|
|
|
|
from tvm import rpc
|
|
|
|
|
|
class EvaluatorConfig(NamedTuple):
|
|
"""Config Details of Evaluator
|
|
|
|
Parameters
|
|
----------
|
|
number: int
|
|
The number of times to run this function for taking average.
|
|
We call these runs as one `repeat` of measurement.
|
|
repeat: int
|
|
The number of times to repeat the measurement.
|
|
In total, the function will be invoked (1 + number x repeat) times,
|
|
where the first one is warm up and will be discarded.
|
|
The returned result contains `repeat` costs,
|
|
each of which is an average of `number` costs.
|
|
min_repeat_ms: int
|
|
Minimum repeat time in ms. if the execution latency is too short,
|
|
increase the number of runs to the given time (in ms) to reduce the measurement error.
|
|
enable_cpu_cache_flush: bool
|
|
Whether to flush the cache on CPU.
|
|
|
|
Note
|
|
----
|
|
The total number of actual executions is 1+number*repeat because we would warm up 1 time before
|
|
actual run. The number of runs would be increased if run time is below min_repeat_ms.
|
|
"""
|
|
|
|
number: int = 3
|
|
repeat: int = 1
|
|
min_repeat_ms: int = 100
|
|
enable_cpu_cache_flush: bool = False
|
|
|
|
@staticmethod
|
|
def _normalized(config: Optional["EvaluatorConfig"]) -> "EvaluatorConfig":
|
|
if config is None:
|
|
return EvaluatorConfig()
|
|
config = EvaluatorConfig(
|
|
number=config.number,
|
|
repeat=config.repeat,
|
|
min_repeat_ms=config.min_repeat_ms,
|
|
enable_cpu_cache_flush=config.enable_cpu_cache_flush,
|
|
)
|
|
return config
|
|
|
|
|
|
class RPCConfig(NamedTuple):
|
|
"""RPC configuration
|
|
|
|
Parameters
|
|
----------
|
|
tracker_host: str
|
|
Host of the RPC Tracker
|
|
tracker_port: int
|
|
Port of the RPC Tracker
|
|
tracker_key: str
|
|
Key of the Tracker
|
|
session_timeout_sec: float
|
|
Timeout of the RPC session
|
|
session_priority: int
|
|
Priority of the RPC session
|
|
"""
|
|
|
|
tracker_host: str | None = None
|
|
tracker_port: None | int | str = None
|
|
tracker_key: str | None = None
|
|
session_priority: int = 1
|
|
session_timeout_sec: int = 10
|
|
|
|
def _sanity_check(self) -> None:
|
|
err_str = (
|
|
"RPCConfig.{0} is not provided. Please provide it explicitly,"
|
|
"or set environment variable {1}"
|
|
)
|
|
if self.tracker_host is None:
|
|
raise ValueError(err_str.format("tracker_host", "TVM_TRACKER_HOST"))
|
|
if self.tracker_port is None:
|
|
raise ValueError(err_str.format("tracker_port", "TVM_TRACKER_PORT"))
|
|
if self.tracker_key is None:
|
|
raise ValueError(err_str.format("tracker_key", "TVM_TRACKER_KEY"))
|
|
|
|
@staticmethod
|
|
def _normalized(config: Optional["RPCConfig"]) -> "RPCConfig":
|
|
if config is None:
|
|
config = RPCConfig()
|
|
tracker_host = config.tracker_host or os.environ.get("TVM_TRACKER_HOST", None)
|
|
tracker_port = config.tracker_port or os.environ.get("TVM_TRACKER_PORT", None)
|
|
tracker_key = config.tracker_key or os.environ.get("TVM_TRACKER_KEY", None)
|
|
if isinstance(tracker_port, str):
|
|
tracker_port = int(tracker_port)
|
|
config = RPCConfig(
|
|
tracker_host=tracker_host,
|
|
tracker_port=tracker_port,
|
|
tracker_key=tracker_key,
|
|
session_priority=config.session_priority,
|
|
session_timeout_sec=config.session_timeout_sec,
|
|
)
|
|
config._sanity_check() # pylint: disable=protected-access
|
|
return config
|
|
|
|
def connect_tracker(self) -> rpc.TrackerSession:
|
|
"""Connect to the tracker
|
|
|
|
Returns
|
|
-------
|
|
tracker : TrackerSession
|
|
The connected tracker session
|
|
"""
|
|
tracker: rpc.TrackerSession | None = None
|
|
|
|
def _connect():
|
|
nonlocal tracker
|
|
tracker = rpc.connect_tracker(self.tracker_host, self.tracker_port)
|
|
|
|
t = Thread(target=_connect)
|
|
t.start()
|
|
t.join(self.session_timeout_sec)
|
|
if t.is_alive() or tracker is None:
|
|
raise ValueError(
|
|
"Unable to connect to the tracker using the following configuration:\n"
|
|
f" tracker host: {self.tracker_host}\n"
|
|
f" tracker port: {self.tracker_port}\n"
|
|
f" timeout (sec): {self.session_timeout_sec}\n"
|
|
"Please check the tracker status via the following command:\n"
|
|
" python3 -m tvm.exec.query_rpc_tracker "
|
|
f"--host {self.tracker_host} --port {self.tracker_port}"
|
|
)
|
|
return tracker
|
|
|
|
def connect_server(self) -> rpc.RPCSession:
|
|
"""Connect to the server
|
|
|
|
Returns
|
|
-------
|
|
session : RPCSession
|
|
The connected rpc session
|
|
"""
|
|
tracker = self.connect_tracker()
|
|
session: rpc.RPCSession = tracker.request(
|
|
key=self.tracker_key,
|
|
priority=self.session_priority,
|
|
session_timeout=self.session_timeout_sec,
|
|
)
|
|
return session
|
|
|
|
def count_num_servers(self, allow_missing=True) -> int:
|
|
"""Count the number of servers available in the tracker
|
|
|
|
Parameters
|
|
----------
|
|
allow_missing : bool
|
|
Whether to allow no server to be found.
|
|
|
|
Returns
|
|
-------
|
|
num_servers : int
|
|
The number of servers
|
|
"""
|
|
tracker = self.connect_tracker()
|
|
tracker_summary = tracker.summary()
|
|
result: int = 0
|
|
for item in tracker_summary["server_info"]:
|
|
_, item_key = item["key"].split(":")
|
|
if item_key == self.tracker_key:
|
|
result += 1
|
|
if result == 0 and not allow_missing:
|
|
raise ValueError(
|
|
"Unable to find servers with the specific key using the following configuration:\n"
|
|
f" tracker host: {self.tracker_host}\n"
|
|
f" tracker port: {self.tracker_port}\n"
|
|
f" tracker key: {self.tracker_key}\n"
|
|
f" timeout (sec): {self.session_timeout_sec}\n"
|
|
"Please check the tracker status via the following command:\n"
|
|
" python3 -m tvm.exec.query_rpc_tracker "
|
|
f"--host {self.tracker_host} --port {self.tracker_port}\n"
|
|
f'and look for key: "{self.tracker_key}"'
|
|
)
|
|
return result
|