Skip to content

Commit 0cecaa5

Browse files
authored
[file report] upload costs to file on s3 (#7050)
I want the pricing information for each runner by calling the API basically caches the data on s3 the file looks like https://ossci-metrics.s3.us-east-1.amazonaws.com/ec2_pricing.json.gz
1 parent 9f2ef5f commit 0cecaa5

File tree

3 files changed

+144
-18
lines changed

3 files changed

+144
-18
lines changed

.github/workflows/update_test_file_ratings.yml

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,3 +104,26 @@ jobs:
104104
user_email: "test-infra@pytorch.org"
105105
user_name: "Pytorch Test Infra"
106106
commit_message: "Updating TD heuristic: historical edited files"
107+
108+
update-ec2-pricing:
109+
runs-on: linux.large
110+
steps:
111+
- name: Checkout test-infra repository
112+
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
113+
114+
- name: Install Dependencies
115+
run: python3 -m pip install boto3==1.19.12 PyYAML==6.0
116+
117+
- name: Generate EC2 pricing data
118+
run: |
119+
python3 tools/torchci/test_insights/ec2_pricing.py
120+
121+
- name: Compress pricing file
122+
run: |
123+
gzip ec2_pricing.json
124+
125+
- name: Upload pricing file to S3
126+
run: |
127+
aws s3 cp ec2_pricing.json.gz s3://ossci-metrics/ec2_pricing.json.gz \
128+
--content-encoding gzip \
129+
--content-type application/json
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
#!/usr/bin/env python3
2+
"""
3+
EC2 Pricing Map Generator
4+
5+
Get pricing info for EC2 instances by reading .github/scale-config.yml and
6+
fetching current AWS pricing data.
7+
"""
8+
9+
import json
10+
from functools import lru_cache
11+
from typing import Optional
12+
13+
import requests
14+
import yaml
15+
16+
17+
@lru_cache
18+
def _get_scale_config() -> dict:
19+
"""Load scale-config.yml and return as a dictionary."""
20+
with open(".github/scale-config.yml", "r") as f:
21+
config = yaml.safe_load(f)
22+
return config
23+
24+
25+
def get_ec2_instance_for_label(label: str) -> dict[str, Optional[str]]:
26+
"""Get EC2 instance type for a given GitHub Actions runner label from scale-config.yml."""
27+
config = _get_scale_config()
28+
29+
runner_info = config.get("runner_types", {})
30+
31+
if label in runner_info:
32+
return {
33+
"ec2_instance": runner_info[label].get("instance_type", None),
34+
"os": runner_info[label].get("os", "linux"),
35+
} # Default to linux if not specified
36+
return {"ec2_instance": None, "os": None}
37+
38+
39+
@lru_cache
40+
def get_all_pricing_data() -> dict:
41+
"""Fetch the entire EC2 pricing data from AWS pricing API. Cached for efficiency."""
42+
price_list_url = "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonEC2/current/us-east-1/index.json"
43+
response = requests.get(price_list_url)
44+
response.raise_for_status()
45+
return response.json()
46+
47+
48+
@lru_cache
49+
def get_price_for_ec2_instance(instance_type, os_type="linux") -> Optional[float]:
50+
"""Fetch on-demand price for EC2 instance type using AWS public pricing data. Returns None if not found."""
51+
52+
# Map os_type to AWS pricing API values
53+
operating_system = "Windows" if os_type.lower() == "windows" else "Linux"
54+
55+
# Get the cached pricing data
56+
pricing_data = get_all_pricing_data()
57+
58+
# Search through the products to find matching instance
59+
for product_sku, product_data in pricing_data.get("products", {}).items():
60+
attributes = product_data.get("attributes", {})
61+
62+
if (
63+
attributes.get("instanceType") == instance_type
64+
and attributes.get("location") == "US East (N. Virginia)"
65+
and attributes.get("operatingSystem") == operating_system
66+
and attributes.get("preInstalledSw") == "NA"
67+
and attributes.get("tenancy") == "Shared"
68+
and attributes.get("usagetype", "").startswith("BoxUsage")
69+
):
70+
# Found the product, now get the pricing terms
71+
terms = (
72+
pricing_data.get("terms", {}).get("OnDemand", {}).get(product_sku, {})
73+
)
74+
75+
for term_data in terms.values():
76+
price_dimensions = term_data.get("priceDimensions", {})
77+
for price_data in price_dimensions.values():
78+
price_per_unit = price_data.get("pricePerUnit", {}).get("USD")
79+
if price_per_unit:
80+
return float(price_per_unit)
81+
82+
print(f"No pricing found for {instance_type} ({operating_system})")
83+
return None
84+
85+
86+
@lru_cache
87+
def get_price_for_label(label: str) -> Optional[float]:
88+
"""Get the on-demand price for the EC2 instance type associated with the given GitHub Actions runner label."""
89+
instance_info = get_ec2_instance_for_label(label)
90+
instance_type = instance_info["ec2_instance"]
91+
os_type = instance_info["os"]
92+
if instance_type is not None:
93+
return get_price_for_ec2_instance(instance_type, os_type)
94+
return None
95+
96+
97+
if __name__ == "__main__":
98+
# Example usage
99+
info = []
100+
scale_config = _get_scale_config()
101+
for runner_label in scale_config.get("runner_types", {}):
102+
price = get_price_for_label(runner_label)
103+
info.append(
104+
{
105+
"label": runner_label,
106+
"price_per_hour": price,
107+
"instance_type": get_ec2_instance_for_label(runner_label)[
108+
"ec2_instance"
109+
],
110+
}
111+
)
112+
with open("ec2_pricing.json", "w") as f:
113+
for line in info:
114+
json.dump(line, f)
115+
f.write("\n")

tools/torchci/test_insights/file_report_generator.py

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
boto3 = None # type: ignore[assignment]
4545

4646
from torchci.clickhouse import query_clickhouse
47+
from torchci.test_insights.ec2_pricing import get_price_for_label
4748

4849

4950
logger = logging.getLogger(__name__)
@@ -74,21 +75,6 @@ def __init__(self, dry_run: bool = True):
7475
"""Initialize the generator with the test owners file path"""
7576
self.dry_run = dry_run
7677

77-
@lru_cache
78-
def load_runner_costs(self) -> Dict[str, float]:
79-
"""Load runner costs from the S3 endpoint"""
80-
logger.debug("Fetching EC2 pricing data from S3...")
81-
with urllib.request.urlopen(self.EC2_PRICING_URL) as response:
82-
compressed_data = response.read()
83-
84-
decompressed_data = gzip.decompress(compressed_data)
85-
pricing_data = {}
86-
for line in decompressed_data.decode("utf-8").splitlines():
87-
if line.strip():
88-
line_json = json.loads(line)
89-
pricing_data[line_json[0]] = float(line_json[2])
90-
return pricing_data
91-
9278
@lru_cache
9379
def load_test_owners(self) -> List[Dict[str, Any]]:
9480
"""Load the test owner labels JSON file from S3"""
@@ -105,10 +91,12 @@ def load_test_owners(self) -> List[Dict[str, Any]]:
10591

10692
def get_runner_cost(self, runner_label: str) -> float:
10793
"""Get the cost per hour for a given runner"""
108-
runner_costs = self.load_runner_costs()
10994
if runner_label.startswith("lf."):
11095
runner_label = runner_label[3:]
111-
return runner_costs.get(runner_label, 0.0)
96+
cost = get_price_for_label(runner_label)
97+
if cost is None:
98+
return 0.0
99+
return cost
112100

113101
def _get_first_suitable_sha(self, shas: list[dict[str, Any]]) -> Optional[str]:
114102
"""Get the first suitable SHA from a list of SHAs."""
@@ -282,7 +270,7 @@ def _get_runner_label_from_job_info(self, job_info: Dict[str, Any]) -> str:
282270
for label in job_labels:
283271
if label.startswith("lf."):
284272
label = label[3:]
285-
if label in self.load_runner_costs():
273+
if get_price_for_label(label) is not None:
286274
return label
287275

288276
return "unknown"

0 commit comments

Comments
 (0)