import subprocess
import sys
import unittest
import pathlib

from torch.testing._internal.common_utils import TestCase, run_tests, IS_LINUX, IS_IN_CI


REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent

try:
    # Just in case PyTorch was not built in 'develop' mode
    sys.path.append(str(REPO_ROOT))
    from tools.stats.scribe import rds_write, register_rds_schema
except ImportError:
    register_rds_schema = None
    rds_write = None


# these tests could eventually be changed to fail if the import/init
# time is greater than a certain threshold, but for now we just use them
# as a way to track the duration of `import torch` in our ossci-metrics
# S3 bucket (see tools/stats/print_test_stats.py)
class TestImportTime(TestCase):
    def test_time_import_torch(self):
        TestCase.runWithPytorchAPIUsageStderr("import torch")

    def test_time_cuda_device_count(self):
        TestCase.runWithPytorchAPIUsageStderr(
            "import torch; torch.cuda.device_count()",
        )

    @unittest.skipIf(not IS_LINUX, "Memory test is only implemented for Linux")
    @unittest.skipIf(not IS_IN_CI, "Memory test only runs in CI")
    def test_peak_memory(self):
        def profile(module, name):
            command = f"import {module}; import resource; print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)"
            result = subprocess.run(
                [sys.executable, "-c", command],
                stdout=subprocess.PIPE,
            )
            max_rss = int(result.stdout.decode().strip())

            return {
                "test_name": name,
                "peak_memory_bytes": max_rss,
            }

        data = profile("torch", "pytorch")
        baseline = profile("sys", "baseline")
        rds_write(
            "import_stats", [data, baseline]
        )


if __name__ == "__main__":
    if register_rds_schema and IS_IN_CI:
        register_rds_schema(
            "import_stats",
            {
                "test_name": "string",
                "peak_memory_bytes": "int",
                "time_ms": "int",
            },
        )

    run_tests()

[((1010, 1084), 'unittest.skipIf', 'unittest.skipIf', (['(not IS_LINUX)', '"""Memory test is only implemented for Linux"""'], {}), "(not IS_LINUX, 'Memory test is only implemented for Linux')\n", (1025, 1084), False, 'import unittest\n'), ((1090, 1150), 'unittest.skipIf', 'unittest.skipIf', (['(not IS_IN_CI)', '"""Memory test only runs in CI"""'], {}), "(not IS_IN_CI, 'Memory test only runs in CI')\n", (1105, 1150), False, 'import unittest\n'), ((2102, 2113), 'torch.testing._internal.common_utils.run_tests', 'run_tests', ([], {}), '()\n', (2111, 2113), False, 'from torch.testing._internal.common_utils import TestCase, run_tests, IS_LINUX, IS_IN_CI\n'), ((794, 847), 'torch.testing._internal.common_utils.TestCase.runWithPytorchAPIUsageStderr', 'TestCase.runWithPytorchAPIUsageStderr', (['"""import torch"""'], {}), "('import torch')\n", (831, 847), False, 'from torch.testing._internal.common_utils import TestCase, run_tests, IS_LINUX, IS_IN_CI\n'), ((900, 985), 'torch.testing._internal.common_utils.TestCase.runWithPytorchAPIUsageStderr', 'TestCase.runWithPytorchAPIUsageStderr', (['"""import torch; torch.cuda.device_count()"""'], {}), "('import torch; torch.cuda.device_count()'\n    )\n", (937, 985), False, 'from torch.testing._internal.common_utils import TestCase, run_tests, IS_LINUX, IS_IN_CI\n'), ((1748, 1791), 'tools.stats.scribe.rds_write', 'rds_write', (['"""import_stats"""', '[data, baseline]'], {}), "('import_stats', [data, baseline])\n", (1757, 1791), False, 'from tools.stats.scribe import rds_write, register_rds_schema\n'), ((1892, 2002), 'tools.stats.scribe.register_rds_schema', 'register_rds_schema', (['"""import_stats"""', "{'test_name': 'string', 'peak_memory_bytes': 'int', 'time_ms': 'int'}"], {}), "('import_stats', {'test_name': 'string',\n    'peak_memory_bytes': 'int', 'time_ms': 'int'})\n", (1911, 2002), False, 'from tools.stats.scribe import rds_write, register_rds_schema\n'), ((1356, 1427), 'subprocess.run', 'subprocess.run', (["[sys.executable, '-c', command]"], {'stdout': 'subprocess.PIPE'}), "([sys.executable, '-c', command], stdout=subprocess.PIPE)\n", (1370, 1427), False, 'import subprocess\n'), ((164, 186), 'pathlib.Path', 'pathlib.Path', (['__file__'], {}), '(__file__)\n', (176, 186), False, 'import pathlib\n')]