wehub-resource-sync
69 lines
1.7 KiB
Python
69 lines
1.7 KiB
Python
import subprocess
|
|
|
|
import requests
|
|
from torch import nn
|
|
|
|
import ray
|
|
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
|
|
|
|
|
class NeuralNetwork(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.flatten = nn.Flatten()
|
|
self.linear_relu_stack = nn.Sequential(
|
|
nn.Linear(28 * 28, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 10),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.flatten(x)
|
|
logits = self.linear_relu_stack(x)
|
|
return logits
|
|
|
|
|
|
def terminate_current_instance():
|
|
"""Use AWS CLI to terminate current instance."""
|
|
|
|
token = requests.put(
|
|
"http://169.254.169.254/latest/api/token",
|
|
headers={"X-aws-ec2-metadata-token-ttl-seconds": "300"},
|
|
timeout=10,
|
|
).text
|
|
instance_id = requests.get(
|
|
"http://169.254.169.254/latest/meta-data/instance-id",
|
|
headers={"X-aws-ec2-metadata-token": token},
|
|
timeout=10,
|
|
).text
|
|
region = requests.get(
|
|
"http://169.254.169.254/latest/meta-data/placement/region",
|
|
headers={"X-aws-ec2-metadata-token": token},
|
|
timeout=10,
|
|
).text
|
|
return subprocess.run(
|
|
[
|
|
"aws",
|
|
"ec2",
|
|
"terminate-instances",
|
|
"--instance-ids",
|
|
instance_id,
|
|
"--region",
|
|
region,
|
|
],
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
|
|
|
|
def terminate_node(node_id: str):
|
|
killer_task = ray.remote(terminate_current_instance).options(
|
|
num_cpus=0,
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(node_id, soft=False),
|
|
)
|
|
ray.get(killer_task.remote())
|