class TestClient:
def setup(self):
self.env_config = get_env_config()
self.trainable_push_groups = self.env_config.trainable_push_groups
self.serve_config = get_serve_config()
self.action_space = ActionSpace(self.trainable_push_groups)
self.action_space.seed(2048)
self.mock_agent = mock_agent_with_action_space(self.action_space)
self.saved_model = SavedModel(self.mock_agent, Environment,
"/tmp/model_location", {})
def test_get_policy(self):
server = Server(self.serve_config)
policy_service = PolicyService(self.saved_model)
server.set_policy_service(policy_service)
try:
server.start()
time.sleep(0.5)
# create the client
address = "{}:{}".format(self.serve_config.host,
self.serve_config.port)
channel = grpc.insecure_channel(address)
client_stub = Client(channel)
policy = client_stub.get_policy(
url="https://www.example.com",
client_env=client.get_random_client_environment(),
manifest=self.env_config,
)
assert policy is not None
assert len(list(policy.push)) + len(list(policy.preload)) > 0
finally:
server.stop()
def setup(self):
self.config = get_config()
self.action_space = ActionSpace(get_push_groups())
self.client_environment = get_random_client_environment()
self.policy = Policy(self.action_space)
applied = True
while applied:
applied = self.policy.apply_action(self.action_space.sample())
def setup(self):
self.env_config = get_env_config()
self.trainable_push_groups = self.env_config.trainable_push_groups
self.serve_config = get_serve_config()
self.action_space = ActionSpace(self.trainable_push_groups)
self.action_space.seed(2048)
self.mock_agent = mock_agent_with_action_space(self.action_space)
self.saved_model = SavedModel(self.mock_agent, Environment,
"/tmp/model_location", {})
def test_observation_with_nonempty_policy_with_default_actions(self):
# use all push groups except the chosen default group
candidate_push_groups = [
i for i, group in enumerate(self.push_groups)
if len(group.resources) > 2 and not group.trainable
]
default_group_idx = random.choice(candidate_push_groups)
default_group = self.push_groups[default_group_idx]
remaining_groups = [
group for i, group in enumerate(self.push_groups)
if i != default_group_idx
]
action_space = ActionSpace(remaining_groups)
policy = Policy(action_space)
# apply some default action
for push in default_group.resources[1:]:
policy.add_default_push_action(default_group.resources[0], push)
# do some actions and check the observation space over time
for _ in range(len(action_space) - 1):
# get an action and apply it in the policy
action_id = action_space.sample()
policy.apply_action(action_id)
# get the observation
obs = get_observation(self.client_environment, self.push_groups,
policy, set())
assert self.observation_space.contains(obs)
# make sure the push sources are recorded correctly
for (source, push) in policy.observable_push:
for push_res in push:
# +1 since we have defined it that way
assert obs["resources"][str(
push_res.order)][-2] == source.source_id + 1
# make sure the push sources are recorded correctly
for (source, preload) in policy.observable_preload:
for push_res in preload:
# +1 since we have defined it that way
assert obs["resources"][str(
push_res.order)][-1] == source.order + 1
# check that all other resources are not pushed
pushed_res = set(push_res.order
for (source, push) in policy.observable_push
for push_res in push)
preloaded_res = set(push_res.order
for (source, push) in policy.observable_preload
for push_res in push)
assert all(res[-2] == 0 for order, res in obs["resources"].items()
if int(order) not in pushed_res)
assert all(res[-1] == 0 for order, res in obs["resources"].items()
if int(order) not in preloaded_res)
def test_init(self):
action_space = ActionSpace(self.env_config.push_groups)
mock_agent = MockAgent(action_space)
m = ModelInstance(mock_agent, self.config)
assert isinstance(m, ModelInstance)
assert m.agent is mock_agent
assert not m._policy
def setup(self):
self.client_environment = get_random_client_environment()
self.page = get_page("http://example.com", self.client_environment)
self.push_groups = get_push_groups()
self.trainable_push_groups = [group for group in self.push_groups if group.trainable]
self.action_space = ActionSpace(self.trainable_push_groups)
self.saved_model = SavedModel(mock_agent_with_action_space(self.action_space), Environment, "", {})
def test_push_policy_returns_cached_policy(self):
action_space = ActionSpace(self.trainable_push_groups)
mock_agent = MockAgent(action_space)
m = ModelInstance(mock_agent, self.config)
first_policy = m.policy
second_policy = m.policy
assert first_policy is second_policy
def test_policy(self):
observation_space = get_observation_space()
action_space = ActionSpace(self.trainable_push_groups)
mock_agent = MockAgent(action_space)
m = ModelInstance(mock_agent, self.config)
policy = m.policy
assert policy
assert all(
observation_space.contains(obs) for obs in mock_agent.observations)
def setup(self):
self.push_groups = get_push_groups()
self.trainable_push_groups = [
group for group in self.push_groups if group.trainable
]
self.serve_config = get_serve_config()
self.action_space = ActionSpace(self.trainable_push_groups)
self.mock_agent = mock_agent_with_action_space(self.action_space)
self.saved_model = SavedModel(self.mock_agent, Environment,
"/tmp/model_location", {})
def test_observation_with_nonempty_policy(self):
action_space = ActionSpace(self.push_groups)
policy = Policy(action_space)
# do some actions and check the observation space over time
for _ in range(len(action_space) - 1):
# get an action and apply it in the policy
action_id = action_space.sample()
policy.apply_action(action_id)
# get the observation
obs = get_observation(self.client_environment, self.push_groups,
policy, set())
assert self.observation_space.contains(obs)
# make sure the push sources are recorded correctly
for (source, push) in policy.push:
for push_res in push:
# +1 since we have defined it that way
assert obs["resources"][str(
push_res.order)][-2] == source.source_id + 1
# make sure the push sources are recorded correctly
for (source, preload) in policy.preload:
for push_res in preload:
# +1 since we have defined it that way
assert obs["resources"][str(
push_res.order)][-1] == source.order + 1
# check that all other resources are not pushed
pushed_res = set(push_res.order for (source, push) in policy.push
for push_res in push)
preloaded_res = set(push_res.order
for (source, push) in policy.preload
for push_res in push)
assert all(res[-2] == 0 for order, res in obs["resources"].items()
if int(order) not in pushed_res)
assert all(res[-1] == 0 for order, res in obs["resources"].items()
if int(order) not in preloaded_res)
def initialize_environment(self,
client_environment: ClientEnvironment,
cached_urls: Optional[Set[str]] = None):
""" Initialize the environment """
log.info(
"initialized environment",
network_type=client.NetworkType(client_environment.network_type),
network_speed=client.NetworkSpeed(
client_environment.network_speed),
device_speed=client.DeviceSpeed(client_environment.device_speed),
bandwidth=client_environment.bandwidth,
latency=client_environment.latency,
cpu_slowdown=client_environment.cpu_slowdown,
loss=client_environment.loss,
reward_func=self.analyzer.reward_func_num,
cached_urls=cached_urls,
)
# Cache scenarios in hours
scenarios = [0, 0, 0, 0, 0, 1, 2, 4, 12, 24]
cache_time = self.np_random.choice(scenarios)
self.cached_urls = (cached_urls if cached_urls is not None else
set() if cache_time == 0 else set(
res.url
for group in self.env_config.push_groups
for res in group.resources
if res.cache_time >= (cache_time * 60 * 60)))
self.client_environment = client_environment
self.analyzer.reset(self.client_environment, self.cached_urls)
num_domains_deployed = math.ceil(PROPORTION_DEPLOYED *
len(self.env_config.push_groups))
push_groups = sorted(self.env_config.push_groups,
key=lambda g: len(g.resources),
reverse=True)[:num_domains_deployed]
self.action_space = ActionSpace(push_groups)
self.policy = Policy(self.action_space)
def test_observation_with_cached_urls(self):
action_space = ActionSpace(self.push_groups)
policy = Policy(action_space)
resources = [
res for group in self.push_groups for res in group.resources
]
mask = [random.randint(0, 2) for _ in range(len(resources))]
cached = [res for (res, include) in zip(resources, mask) if include]
cached_urls = set(res.url for res in cached)
obs = get_observation(self.client_environment, self.push_groups,
policy, cached_urls)
for res in cached:
assert obs["resources"][str(res.order)][1] == 1
class TestMahiMahiConfig:
def setup(self):
self.config = get_config()
self.action_space = ActionSpace(get_push_groups())
self.client_environment = get_random_client_environment()
self.policy = Policy(self.action_space)
applied = True
while applied:
applied = self.policy.apply_action(self.action_space.sample())
def test_init_without_policy(self):
mm_config = MahiMahiConfig(self.config)
assert isinstance(mm_config, MahiMahiConfig)
assert mm_config.policy is None
assert mm_config.client_environment is None
def test_init_without_client_environment(self):
mm_config = MahiMahiConfig(self.config, policy=self.policy)
assert isinstance(mm_config, MahiMahiConfig)
assert mm_config.policy is self.policy
assert mm_config.client_environment is None
def test_init_with_client_environment(self):
mm_config = MahiMahiConfig(self.config, policy=self.policy, client_environment=self.client_environment)
assert isinstance(mm_config, MahiMahiConfig)
assert mm_config.policy is self.policy
assert mm_config.client_environment is self.client_environment
def test_record_shell_with_cmd(self):
save_dir = "/tmp/save_dir"
mm_config = MahiMahiConfig(self.config, policy=self.policy)
cmd = mm_config.record_shell_with_cmd(save_dir, ["a", "command"])
assert cmd == (mm_config.record_cmd(save_dir) + ["a", "command"])
def test_record_cmd(self):
save_dir = "/tmp/save_dir"
mm_config = MahiMahiConfig(self.config)
record_cmd = mm_config.record_cmd(save_dir)
assert record_cmd[0] == "mm-webrecord"
assert record_cmd[1] == save_dir
def test_formatted_trace_file(self):
mm_config = MahiMahiConfig(self.config, policy=self.policy, client_environment=self.client_environment)
trace_lines = trace_for_kbps(self.client_environment.bandwidth)
formatted = format_trace_lines(trace_lines)
assert mm_config.formatted_trace_file == formatted
def test_get_default_observation(self):
action_space = ActionSpace(self.push_groups)
policy = Policy(action_space)
obs = get_observation(self.client_environment, self.push_groups,
policy, set())
assert isinstance(obs, dict)
assert self.observation_space.contains(obs)
# assert that the client environment is correctly captured
assert obs["client"][
"network_type"] == self.client_environment.network_type.value
assert obs["client"][
"device_speed"] == self.client_environment.device_speed.value
# assert that all resources are not pushed initially
assert all(res[-2] == 0 for res in obs["resources"].values())
# assert that all resources are not preloaded initially
assert all(res[-1] == 0 for res in obs["resources"].values())
# assert that the push_groups are encoded correctly
for group in self.push_groups:
for res in group.resources:
assert np.array_equal(
obs["resources"][str(res.order)],
np.array((
1, # resource is enabled
0, # resource is not cached
group.id, # the resource's domain id
res.
source_id, # the resource's relative offset from its domain top
res.order +
1, # the resource's absolute offset from the start of the page load
res.initiator + 1, # the resource's initiator
res.type.value, # resource type
res.size // 1000, # resource size in KB
0, # not pushed
0, # not preloaded
)),
)
max_order = max(r.order for group in self.push_groups
for r in group.resources)
for i in range(max_order + 1, MAX_RESOURCES):
assert np.array_equal(obs["resources"][str(i)],
np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]))
def get_mahimahi_config() -> MahiMahiConfig:
return MahiMahiConfig(
config=get_config(),
policy=Policy(ActionSpace(get_push_groups())),
client_environment=get_random_client_environment(),
)
def setup(self):
self.config = get_config()
self.policy = Policy(ActionSpace(self.config.env_config.push_groups))
self.client_environment = get_random_client_environment()
def get_action(action_space: ActionSpace) -> ActionIDType:
# pick a non-noop action
action = NOOP_ACTION_ID
while action_space.decode_action(action).is_noop:
action = action_space.sample()
return action
class Environment(gym.Env):
"""
Environment virtualizes a randomly chosen network and browser environment and
facilitates the training for a given web page. This includes action selection, policy
generation, and evaluation of the policy/action in the simulated environment.
"""
def __init__(self, config: Union[Config, dict]):
# make sure config is an instance of Config or a dict
assert isinstance(config, (Config, dict))
config = config if isinstance(config, Config) else Config(**config)
self.config = config
self.env_config = config.env_config
self.np_random = np.random.RandomState()
log.info("initialized trainable push groups",
groups=[
group.name
for group in self.env_config.trainable_push_groups
])
self.observation_space = get_observation_space()
self.cached_urls = config.cached_urls or set()
self.analyzer = Analyzer(self.config, config.reward_func or 0,
config.use_aft or False)
self.client_environment: Optional[ClientEnvironment] = None
self.action_space: Optional[ActionSpace] = None
self.policy: Optional[Policy] = None
self.initialize_environment(
self.config.client_env
or client.get_random_fast_lte_client_environment(),
self.config.cached_urls)
def seed(self, seed=None):
self.np_random.seed(seed)
def reset(self):
self.initialize_environment(
client.get_random_fast_lte_client_environment(),
self.config.cached_urls)
return self.observation
def initialize_environment(self,
client_environment: ClientEnvironment,
cached_urls: Optional[Set[str]] = None):
""" Initialize the environment """
log.info(
"initialized environment",
network_type=client.NetworkType(client_environment.network_type),
network_speed=client.NetworkSpeed(
client_environment.network_speed),
device_speed=client.DeviceSpeed(client_environment.device_speed),
bandwidth=client_environment.bandwidth,
latency=client_environment.latency,
cpu_slowdown=client_environment.cpu_slowdown,
loss=client_environment.loss,
reward_func=self.analyzer.reward_func_num,
cached_urls=cached_urls,
)
# Cache scenarios in hours
scenarios = [0, 0, 0, 0, 0, 1, 2, 4, 12, 24]
cache_time = self.np_random.choice(scenarios)
self.cached_urls = (cached_urls if cached_urls is not None else
set() if cache_time == 0 else set(
res.url
for group in self.env_config.push_groups
for res in group.resources
if res.cache_time >= (cache_time * 60 * 60)))
self.client_environment = client_environment
self.analyzer.reset(self.client_environment, self.cached_urls)
num_domains_deployed = math.ceil(PROPORTION_DEPLOYED *
len(self.env_config.push_groups))
push_groups = sorted(self.env_config.push_groups,
key=lambda g: len(g.resources),
reverse=True)[:num_domains_deployed]
self.action_space = ActionSpace(push_groups)
self.policy = Policy(self.action_space)
def step(self, action: ActionIDType):
# decode the action and apply it to the policy
decoded_action = self.action_space.decode_action(action)
action_applied = self.policy.apply_action(decoded_action)
# make sure the action isn't used again
log.info("trying action",
action_id=action,
action=repr(decoded_action),
steps_taken=self.policy.steps_taken)
self.action_space.use_action(decoded_action)
reward = NOOP_ACTION_REWARD
if action_applied:
reward = self.analyzer.get_reward(self.policy)
log.info("got reward", action=repr(decoded_action), reward=reward)
info = {"action": decoded_action, "policy": self.policy.as_dict}
return self.observation, reward, not action_applied, info
def render(self, mode="human"):
return super(Environment, self).render(mode=mode)
@property
def observation(self):
""" Returns an observation for the current state of the environment """
return get_observation(self.client_environment,
self.env_config.push_groups, self.policy,
self.cached_urls)
def get_action_space():
action_space = ActionSpace(get_push_groups())
action_space.seed(4096)
return action_space
