Context Manager API Enhancement (#159)

Summary:
## Overview
This PR enhances the `TritonParseManager` context manager by adding support for tensor blob storage configuration.

## Changes
- **New Parameters Added to `TritonParseManager.__init__`**:
  - `enable_tensor_blob_storage` (bool, default: `False`): Enables tensor blob storage when set to `True`
  - `tensor_storage_quota` (int, optional, default: `None`): Sets storage quota in bytes for tensor blobs (default quota is 100GB when enabled)

- **Implementation Details**:
  - The new parameters are properly stored as instance attributes
  - Parameters are passed through to the `init()` function in the `__enter__` method
  - Uses conditional inclusion: `tensor_storage_quota` is only passed to `init()` when explicitly set (not `None`)

## Impact
- Backward compatible: All new parameters have default values
- Provides users with fine-grained control over tensor blob storage behavior
- Maintains clean API design by conditionally passing parameters only when needed

Pull Request resolved: #159

Reviewed By: sfzhu93

Differential Revision: D84107701

Pulled By: FindHao

fbshipit-source-id: d8918a948f7c14a820e5a701c1ef15079e30583d

Author: FindHao

tritonparse/context_manager.py

@@ -17,6 +17,8 @@ def __init__(
         self,
         enable_trace_launch=False,
         split_inductor_compilations=True,
+        enable_tensor_blob_storage=False,
+        tensor_storage_quota=None,
         **parse_kwargs,
     ):
         """
@@ -25,17 +27,28 @@ def __init__(
         Args:
             enable_trace_launch: Whether to enable trace launch
             split_inductor_compilations: Whether to split inductor compilations in the output
+            enable_tensor_blob_storage: Whether to enable tensor blob storage
+            tensor_storage_quota: Storage quota in bytes for tensor blobs (default: 100GB)
             **parse_kwargs: Additional keyword arguments to pass to unified_parse
         """
         self.enable_trace_launch = enable_trace_launch
         self.split_inductor_compilations = split_inductor_compilations
+        self.enable_tensor_blob_storage = enable_tensor_blob_storage
+        self.tensor_storage_quota = tensor_storage_quota
         self.parse_kwargs = parse_kwargs
         self.dir_path = None
         self.output_link = None
 
     def __enter__(self):
         self.dir_path = createUniqueTempDirectory()
-        init(self.dir_path, enable_trace_launch=self.enable_trace_launch)
+        init_kwargs = {
+            "enable_trace_launch": self.enable_trace_launch,
+            "enable_tensor_blob_storage": self.enable_tensor_blob_storage,
+        }
+        if self.tensor_storage_quota is not None:
+            init_kwargs["tensor_storage_quota"] = self.tensor_storage_quota
+
+        init(self.dir_path, **init_kwargs)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):