mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
99f2491af9
This reverts commit 45411d1fc9a2b6d2f891b6ab0ae16409719e09fc. Reverted https://github.com/pytorch/pytorch/pull/129409 on behalf of https://github.com/jeanschmidt due to Breaking internal CI, @albanD please help get this PR merged ([comment](https://github.com/pytorch/pytorch/pull/129409#issuecomment-2571316444))
284 lines
8.2 KiB
Python
284 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import gzip
|
|
import inspect
|
|
import json
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest import mock
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from tools.stats.upload_metrics import add_global_metric, emit_metric, global_metrics
|
|
from tools.stats.upload_stats_lib import get_s3_resource, remove_nan_inf
|
|
|
|
|
|
sys.path.remove(str(REPO_ROOT))
|
|
|
|
# default values
|
|
REPO = "some/repo"
|
|
BUILD_ENV = "cuda-10.2"
|
|
TEST_CONFIG = "test-config"
|
|
WORKFLOW = "some-workflow"
|
|
JOB = "some-job"
|
|
RUN_ID = 56
|
|
RUN_NUMBER = 123
|
|
RUN_ATTEMPT = 3
|
|
PR_NUMBER = 6789
|
|
JOB_ID = 234
|
|
JOB_NAME = "some-job-name"
|
|
|
|
|
|
@mock.patch("boto3.resource")
|
|
class TestUploadStats(unittest.TestCase):
|
|
emitted_metric: dict[str, Any] = {"did_not_emit": True}
|
|
|
|
def mock_put_item(self, **kwargs: Any) -> None:
|
|
# Utility for mocking putting items into s3. THis will save the emitted
|
|
# metric so tests can check it
|
|
self.emitted_metric = json.loads(
|
|
gzip.decompress(kwargs["Body"]).decode("utf-8")
|
|
)
|
|
|
|
# Before each test, set the env vars to their default values
|
|
def setUp(self) -> None:
|
|
get_s3_resource.cache_clear()
|
|
global_metrics.clear()
|
|
|
|
mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"CI": "true",
|
|
"BUILD_ENVIRONMENT": BUILD_ENV,
|
|
"TEST_CONFIG": TEST_CONFIG,
|
|
"GITHUB_REPOSITORY": REPO,
|
|
"GITHUB_WORKFLOW": WORKFLOW,
|
|
"GITHUB_JOB": JOB,
|
|
"GITHUB_RUN_ID": str(RUN_ID),
|
|
"GITHUB_RUN_NUMBER": str(RUN_NUMBER),
|
|
"GITHUB_RUN_ATTEMPT": str(RUN_ATTEMPT),
|
|
"JOB_ID": str(JOB_ID),
|
|
"JOB_NAME": str(JOB_NAME),
|
|
},
|
|
clear=True, # Don't read any preset env vars
|
|
).start()
|
|
|
|
def test_emits_default_and_given_metrics(self, mock_resource: Any) -> None:
|
|
metric = {
|
|
"some_number": 123,
|
|
"float_number": 32.34,
|
|
}
|
|
|
|
# Querying for this instead of hard coding it b/c this will change
|
|
# based on whether we run this test directly from python or from
|
|
# pytest
|
|
current_module = inspect.getmodule(inspect.currentframe()).__name__ # type: ignore[union-attr]
|
|
|
|
emit_should_include = {
|
|
"metric_name": "metric_name",
|
|
"calling_file": "test_upload_stats_lib.py",
|
|
"calling_module": current_module,
|
|
"calling_function": "test_emits_default_and_given_metrics",
|
|
"repo": REPO,
|
|
"workflow": WORKFLOW,
|
|
"build_environment": BUILD_ENV,
|
|
"job": JOB,
|
|
"test_config": TEST_CONFIG,
|
|
"run_id": RUN_ID,
|
|
"run_number": RUN_NUMBER,
|
|
"run_attempt": RUN_ATTEMPT,
|
|
"job_id": JOB_ID,
|
|
"job_name": JOB_NAME,
|
|
"info": metric,
|
|
}
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertEqual(
|
|
self.emitted_metric,
|
|
{**self.emitted_metric, **emit_should_include},
|
|
)
|
|
|
|
def test_when_global_metric_specified_then_it_emits_it(
|
|
self, mock_resource: Any
|
|
) -> None:
|
|
metric = {
|
|
"some_number": 123,
|
|
}
|
|
|
|
global_metric_name = "global_metric"
|
|
global_metric_value = "global_value"
|
|
|
|
add_global_metric(global_metric_name, global_metric_value)
|
|
|
|
emit_should_include = {
|
|
**metric,
|
|
global_metric_name: global_metric_value,
|
|
}
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertEqual(
|
|
self.emitted_metric,
|
|
{**self.emitted_metric, "info": emit_should_include},
|
|
)
|
|
|
|
def test_when_local_and_global_metric_specified_then_global_is_overridden(
|
|
self, mock_resource: Any
|
|
) -> None:
|
|
global_metric_name = "global_metric"
|
|
global_metric_value = "global_value"
|
|
local_override = "local_override"
|
|
|
|
add_global_metric(global_metric_name, global_metric_value)
|
|
|
|
metric = {
|
|
"some_number": 123,
|
|
global_metric_name: local_override,
|
|
}
|
|
|
|
emit_should_include = {
|
|
**metric,
|
|
global_metric_name: local_override,
|
|
}
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertEqual(
|
|
self.emitted_metric,
|
|
{**self.emitted_metric, "info": emit_should_include},
|
|
)
|
|
|
|
def test_when_optional_envvar_set_to_actual_value_then_emit_vars_emits_it(
|
|
self, mock_resource: Any
|
|
) -> None:
|
|
metric = {
|
|
"some_number": 123,
|
|
}
|
|
|
|
emit_should_include = {
|
|
"info": {**metric},
|
|
"pr_number": PR_NUMBER,
|
|
}
|
|
|
|
mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"PR_NUMBER": str(PR_NUMBER),
|
|
},
|
|
).start()
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertEqual(
|
|
self.emitted_metric,
|
|
{**self.emitted_metric, **emit_should_include},
|
|
)
|
|
|
|
def test_when_optional_envvar_set_to_a_empty_str_then_emit_vars_ignores_it(
|
|
self, mock_resource: Any
|
|
) -> None:
|
|
metric = {"some_number": 123}
|
|
|
|
emit_should_include: dict[str, Any] = metric.copy()
|
|
|
|
# Github Actions defaults some env vars to an empty string
|
|
default_val = ""
|
|
mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"PR_NUMBER": default_val,
|
|
},
|
|
).start()
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertEqual(
|
|
self.emitted_metric,
|
|
{**self.emitted_metric, "info": emit_should_include},
|
|
f"Metrics should be emitted when an option parameter is set to '{default_val}'",
|
|
)
|
|
self.assertFalse(
|
|
self.emitted_metric.get("pr_number"),
|
|
f"Metrics should not include optional item 'pr_number' when it's envvar is set to '{default_val}'",
|
|
)
|
|
|
|
def test_no_metrics_emitted_if_required_env_var_not_set(
|
|
self, mock_resource: Any
|
|
) -> None:
|
|
metric = {"some_number": 123}
|
|
|
|
mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"CI": "true",
|
|
"BUILD_ENVIRONMENT": BUILD_ENV,
|
|
},
|
|
clear=True,
|
|
).start()
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertTrue(self.emitted_metric["did_not_emit"])
|
|
|
|
def test_no_metrics_emitted_if_required_env_var_set_to_empty_string(
|
|
self, mock_resource: Any
|
|
) -> None:
|
|
metric = {"some_number": 123}
|
|
|
|
mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"GITHUB_JOB": "",
|
|
},
|
|
).start()
|
|
|
|
mock_resource.return_value.Object.return_value.put = self.mock_put_item
|
|
|
|
emit_metric("metric_name", metric)
|
|
|
|
self.assertTrue(self.emitted_metric["did_not_emit"])
|
|
|
|
def test_remove_nan_inf(self, _mocked_resource: Any) -> None:
|
|
checks = [
|
|
(float("inf"), '"inf"', "Infinity"),
|
|
(float("nan"), '"nan"', "NaN"),
|
|
({1: float("inf")}, '{"1": "inf"}', '{"1": Infinity}'),
|
|
([float("nan")], '["nan"]', "[NaN]"),
|
|
({1: [float("nan")]}, '{"1": ["nan"]}', '{"1": [NaN]}'),
|
|
]
|
|
|
|
for input, clean, unclean in checks:
|
|
clean_output = json.dumps(remove_nan_inf(input))
|
|
unclean_output = json.dumps(input)
|
|
self.assertEqual(
|
|
clean_output,
|
|
clean,
|
|
f"Expected {clean} when input is {unclean}, got {clean_output}",
|
|
)
|
|
self.assertEqual(
|
|
unclean_output,
|
|
unclean,
|
|
f"Expected {unclean} when input is {unclean}, got {unclean_output}",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|