Add S3 uploads for profiles

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)
+```

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

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

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
 

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}")