Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,31 @@
import pickle
import re
import sys
import sysconfig
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
# cProfile imports `profile` by name; expose the stdlib API when this CLI shadows it.
if __name__ == "profile":
stdlib_dir = sysconfig.get_path("stdlib")
if not stdlib_dir:
raise ImportError("Could not determine the Python standard library path")
stdlib_path = Path(stdlib_dir) / "profile.py"
spec = importlib.util.spec_from_file_location("_stdlib_profile", stdlib_path)
if spec is None or spec.loader is None:
raise ImportError(f"Could not load stdlib profile module from {stdlib_path}")
stdlib_profile = importlib.util.module_from_spec(spec)
spec.loader.exec_module(stdlib_profile)
__all__ = stdlib_profile.__all__
_Utils = stdlib_profile._Utils
Profile = stdlib_profile.Profile
run = stdlib_profile.run
runctx = stdlib_profile.runctx
else:
import torch
import torch.nn as nn

# ---------------------------------------------------------------------------
# Constants
Expand Down
38 changes: 38 additions & 0 deletions tests/test_profile_module_compat.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
import subprocess
import sys
import unittest
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]


class ProfileModuleCompatibilityTest(unittest.TestCase):
def run_python(self, code):
return subprocess.run(
[sys.executable, "-c", code],
cwd=REPO_ROOT,
capture_output=True,
text=True,
)

def test_profile_exposes_stdlib_api_without_project_imports(self):
result = self.run_python(
"import sys, profile; "
"assert callable(profile.run); "
"assert callable(profile.runctx); "
"assert profile.Profile; "
"assert profile._Utils; "
"assert 'torch' not in sys.modules"
)

self.assertEqual(result.returncode, 0, result.stderr)

def test_cprofile_runs_from_repo_root(self):
result = self.run_python("import cProfile; cProfile.run('sum(range(10))')")

self.assertEqual(result.returncode, 0, result.stderr)


if __name__ == "__main__":
unittest.main()