diff --git a/README.md b/README.md index e5be69b..8a01e23 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ you are really not sure why. Profiling helps to find and analyze these bottlenec into fixing performance problems. Uses [snakeviz](https://jiffyclub.github.io/snakeviz/) for interactive visualizations. Install it with `pip install profiling-helpers`. +Visualize profile files with `snakeviz profile_xyz.prof`. There are two decorators, `time_it` and `profile_it`. Use them anywhere in your code, like this: @@ -32,3 +33,18 @@ def my_slow_function(x): my_slow_function(42) # Opens snakeviz after this function is completed ``` + +Profiles are normally saved on the local file system. If you have other save targets, you can +either use included FileSavers (currently only for AWS S3) or implement your own one by inheriting +from the `BaseFileSaver` class. Here is a variant with S3: + +```python +from profiling_helpers import profile_it, S3FileSaver + +@profile_it(S3FileSaver("s3://my-bucket/my/path/to/profiles/", kms_key_id="...")) +def my_slow_function(x): + sleep(10) + return x + +my_slow_function(42) +``` diff --git a/setup.cfg b/setup.cfg index 1beebe0..03a0757 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = profiling-helpers -version = 0.1.0 +version = 0.2.0 author = Fabian Nobis author_email = fabiannobis@gmx.de url = https://github.com/NobisIndustries/python-profiling-helpers @@ -12,7 +12,7 @@ license_file = LICENSE platform = any keywords = cProfile, profiling, performance classifiers = - Development Status :: 3 - Alpha + Development Status :: 4 - Beta Intended Audience :: Developers License :: OSI Approved :: MIT License Operating System :: OS Independent @@ -29,6 +29,8 @@ install_requires = snakeviz>=2.1.1 [options.extras_require] +aws = + boto3>=1.28.0 dev = pre-commit build diff --git a/src/profiling_helpers/__init__.py b/src/profiling_helpers/__init__.py index d7adf76..9a7b514 100644 --- a/src/profiling_helpers/__init__.py +++ b/src/profiling_helpers/__init__.py @@ -1 +1,2 @@ from profiling_helpers.profiling import profile_it, time_it # noqa: F401 +from profiling_helpers.save_targets import BaseFileSaver, S3FileSaver # noqa: F401 diff --git a/src/profiling_helpers/profiling.py b/src/profiling_helpers/profiling.py index caa770d..a34e82b 100644 --- a/src/profiling_helpers/profiling.py +++ b/src/profiling_helpers/profiling.py @@ -4,8 +4,11 @@ from cProfile import Profile from datetime import datetime from pathlib import Path +from tempfile import NamedTemporaryFile from typing import Union +from profiling_helpers.save_targets import BaseFileSaver, LocalFileSaver + def time_it(function): """ @@ -28,7 +31,9 @@ def inner(*args, **kwargs): return inner -def profile_it(save_dir: Union[str, Path], open_visualization: bool = False): +def profile_it( + save_at: Union[str, Path, BaseFileSaver], open_visualization: bool = False +): """ Provides detailed inspection of function performance with the help of cProfile. The profile files are saved into the given directory and can be explored with snakeviz. @@ -37,21 +42,27 @@ def profile_it(save_dir: Union[str, Path], open_visualization: bool = False): @profile_it('/my/profile/dir') def my_function(x, y, z): ... - :param save_dir: Directory where the profile file is saved. + You can also save your profiles to other places than the local file system. To do this, give save_at an initialized + FileSaver object instead of a local path. For example, you can upload it automatically to AWS S3 like this: + + @profile_it(S3FileSaver("s3://my-bucket/my/path/to/profiles/")) + def my_function(x, y, z): ... + + :param save_at: Directory where the profile file is saved (if you want to save it locally) or a file saver object + that implements the BaseFileServer class if you want to save it in other places, like an S3 bucket. :param open_visualization: If true, automatically open the profile in snakeviz visualizer afterwards. This will - block the program flow after this function call. You can also open it manually with `snakeviz profile_xyz.prof`. + block the program flow after this function call and will only work if you saved your profiles locally. + You can also open profiles manually with `snakeviz profile_xyz.prof`. """ - save_dir = Path(save_dir) - save_dir.mkdir(exist_ok=True, parents=True) + if isinstance(save_at, BaseFileSaver): + file_saver = save_at + else: + file_saver = LocalFileSaver(save_at) def decorator(function): def inner(*args, **kwargs): profiler = Profile() - file_path = Path( - save_dir, - f"{function.__name__}_{datetime.now().isoformat().replace(':', '_').replace('.', '_')}.prof", - ) - + file_name = f"{function.__name__}_{datetime.now().isoformat().replace(':', '_').replace('.', '_')}.prof" try: result = profiler.runcall(function, *args, **kwargs) return result @@ -60,10 +71,18 @@ def inner(*args, **kwargs): raise finally: stats = pstats.Stats(profiler) - stats.dump_stats(file_path.as_posix()) - print(f"Profile saved at {file_path}") + + with NamedTemporaryFile(delete=False) as temp: + stats.dump_stats(temp.name) + with open(temp.name, "rb") as f: + stats_bytes = f.read() + file_saver.save_profile(stats_bytes, file_name) if open_visualization: - os.system(f'snakeviz "{file_path.as_posix()}"') + if not isinstance(file_saver, LocalFileSaver): + raise RuntimeError( + "You need to save your profile locally for automatic snakeviz invocation." + ) + os.system(f'snakeviz "{Path(save_at, file_name).as_posix()}"') return inner diff --git a/src/profiling_helpers/save_targets.py b/src/profiling_helpers/save_targets.py new file mode 100644 index 0000000..e4274b8 --- /dev/null +++ b/src/profiling_helpers/save_targets.py @@ -0,0 +1,75 @@ +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Optional, Tuple, Union + + +class BaseFileSaver(ABC): + @abstractmethod + def save_profile(self, profile_binary: bytes, profile_file_name: str) -> None: + """ + This method gets the bytes content of the profile file you want to save somewhere + and the suggested file name. You can implement your own logic in a child class that + saves the profile on the target you want. + """ + raise NotImplementedError() + + +class LocalFileSaver(BaseFileSaver): + def __init__(self, save_dir: Union[str, Path]): + """ + Saves a file to a local dir. + """ + save_dir = Path(save_dir) + save_dir.mkdir(exist_ok=True, parents=True) + self._save_dir = save_dir + + def save_profile(self, profile_binary: bytes, profile_file_name: str) -> None: + save_path = Path(self._save_dir, profile_file_name) + save_path.write_bytes(profile_binary) + print(f"Profile saved at {save_path.as_posix()}") + + +class S3FileSaver(BaseFileSaver): + def __init__( + self, s3_save_dir: str, kms_key_id: Optional[str] = None, **kwargs + ) -> None: + """ + Uploads profile files to an S3 bucket. Additional keyword arguments are given to the boto3 + S3 client initialization. + + :param s3_save_dir: The base path of the file you want to upload, e.g. "s3://my-bucket/my/path/to/profiles/" + :param kms_key_id: If your target bucket is encrypted with a KMS key, you can provide its ID here. + """ + super().__init__() + try: + import boto3 + except ImportError: + raise ImportError( + "You do not have the optional dependencies installed to upload your profiles to an " + "S3 bucket. Please install these with `pip install profiling-helpers[aws]`" + ) + + self._s3 = boto3.client("s3", **kwargs) + self._s3_bucket, self._s3_prefix = self._split_s3_path(s3_save_dir) + self._kms_key_id = kms_key_id + + @classmethod + def _split_s3_path(cls, s3_path: str) -> Tuple[str, str]: + # Normally we would use .removeprefix(), but we support Python 3.7 which doesn't have that yet + if s3_path.startswith("s3://"): + s3_path = s3_path[5:] + s3_path = s3_path.rstrip("/") + bucket_name, prefix = s3_path.split("/", 1) + return bucket_name, prefix + + def save_profile(self, profile_binary: bytes, profile_file_name: str) -> None: + save_key = f"{self._s3_prefix}/{profile_file_name}" + extra_args = {} + if self._kms_key_id: + extra_args.update( + {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": self._kms_key_id} + ) + self._s3.put_object( + Body=profile_binary, Bucket=self._s3_bucket, Key=save_key, **extra_args + ) + print(f"Profile uploaded to s3://{self._s3_bucket}/{save_key}")