chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,201 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user