Add S3 export functionality with database support

Author: Khauneesh-AI

alembic/versions/1a8fdc23eb6f_add_s3_export_path.py

@@ -0,0 +1,38 @@
+"""add_s3_export_path
+
+Revision ID: 1a8fdc23eb6f
+Revises: 9023b46c8d4c
+Create Date: 2025-04-22 20:01:13.247491
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '1a8fdc23eb6f'
+down_revision: Union[str, None] = '9023b46c8d4c'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Add s3_export_path column to generation_metadata table
+    with op.batch_alter_table('generation_metadata', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('s3_export_path', sa.Text(), nullable=True))
+
+    # Add s3_export_path column to export_metadata table
+    with op.batch_alter_table('export_metadata', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('s3_export_path', sa.Text(), nullable=True))
+
+
+def downgrade() -> None:
+    # Remove s3_export_path column from generation_metadata table
+    with op.batch_alter_table('generation_metadata', schema=None) as batch_op:
+        batch_op.drop_column('s3_export_path')
+
+    # Remove s3_export_path column from export_metadata table
+    with op.batch_alter_table('export_metadata', schema=None) as batch_op:
+        batch_op.drop_column('s3_export_path')

app/core/database.py

@@ -60,6 +60,7 @@ def init_db(self):
                         display_name TEXT,
                         local_export_path TEXT,
                         hf_export_path TEXT,
+                        s3_export_path TEXT,
                         num_questions FLOAT,
                         total_count FLOAT,
                         topics TEXT,
@@ -107,6 +108,7 @@ def init_db(self):
                         display_name TEXT,
                         local_export_path TEXT,
                         hf_export_path TEXT,
+                        s3_export_path TEXT,
                         job_id TEXT,
                         job_name TEXT UNIQUE,
                         job_status TEXT,
@@ -145,29 +147,24 @@ def save_generation_metadata(self, metadata: Dict) -> int:
         try:
             # Prepare data outside transaction
             if metadata.get('generate_file_name'):
-                
                 output_paths = metadata.get('output_path', {})
             else:
-                
                 output_paths = {}
-                
-            
 
             # Use a single connection with enhanced settings
             with self.get_connection() as conn:
                 conn.execute("BEGIN IMMEDIATE")
 
-                
                 cursor = conn.cursor()
 
                 query = """
                 INSERT INTO generation_metadata (
-                timestamp, technique, model_id, inference_type,  caii_endpoint, use_case,
-                custom_prompt, model_parameters, input_key, output_key, output_value, generate_file_name,
-                display_name, local_export_path, hf_export_path,
-                num_questions, total_count, topics, examples, 
-                schema, doc_paths, input_path, job_id, job_name, job_status, job_creator_name
-            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    timestamp, technique, model_id, inference_type, caii_endpoint, use_case,
+                    custom_prompt, model_parameters, input_key, output_key, output_value, generate_file_name,
+                    display_name, local_export_path, hf_export_path, s3_export_path,
+                    num_questions, total_count, topics, examples, 
+                    schema, doc_paths, input_path, job_id, job_name, job_status, job_creator_name
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """
 
                 values = (
@@ -186,6 +183,7 @@ def save_generation_metadata(self, metadata: Dict) -> int:
                     metadata.get('display_name', None),
                     output_paths.get('local', None),
                     output_paths.get('huggingface', None),
+                    output_paths.get('s3', None),
                     metadata.get('num_questions', None),
                     metadata.get('total_count', None),
                     metadata.get('topics', None),
@@ -198,20 +196,18 @@ def save_generation_metadata(self, metadata: Dict) -> int:
                     metadata.get('job_status', None),
                     metadata.get('job_creator_name', None)
                 )
-                #print(values)
 
                 cursor.execute(query, values)
                 conn.commit()
                 return cursor.lastrowid
 
         except sqlite3.OperationalError as e:
-            if conn:
+            if 'conn' in locals():
                 conn.rollback()
             print(f"Database operation error in save_generation_metadata: {e}")
-            
             raise
         except Exception as e:
-            if conn:
+            if 'conn' in locals():
                 conn.rollback()
             print(f"Error saving metadata to database: {str(e)}")
             raise
@@ -359,7 +355,6 @@ def save_evaluation_metadata(self, metadata: Dict) -> int:
     def save_export_metadata(self, metadata: Dict) -> int:
         """Save export metadata to database with prepared transaction"""
         try:
-          
             # Use a single connection with enhanced settings
             with self.get_connection() as conn:
                 conn.execute("BEGIN IMMEDIATE")
@@ -373,11 +368,12 @@ def save_export_metadata(self, metadata: Dict) -> int:
                         display_name,
                         local_export_path,
                         hf_export_path,
+                        s3_export_path,
                         job_id,
                         job_name,
                         job_status,
                         job_creator_name
-                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """
 
                 values = (
@@ -386,6 +382,7 @@ def save_export_metadata(self, metadata: Dict) -> int:
                     metadata.get('display_name'),
                     metadata.get('local_export_path', None),
                     metadata.get('hf_export_path', None),
+                    metadata.get('s3_export_path', None),  # Add this line
                     metadata.get('job_id', None),
                     metadata.get('job_name', None),
                     metadata.get('job_status', None),
@@ -1131,4 +1128,21 @@ def backup_and_restore_db(self, force_restore: bool = False) -> bool:
                     print(f"Force restore failed: {str(restore_error)}")
             return False
 
+    def update_s3_path(self, file_name: str, s3_path: str):
+        """Update s3_export_path for a generation"""
+        try:
+            with self.get_connection() as conn:
+                conn.execute("BEGIN IMMEDIATE")
+                cursor = conn.cursor()
+                
+                # Update the s3_path
+                cursor.execute(
+                    "UPDATE generation_metadata SET s3_export_path = ? WHERE generate_file_name = ?",
+                    (s3_path, file_name)
+                )
+                conn.commit()
+                print(f"S3 path update successful for file: {file_name}")
+        except Exception as e:
+            print(f"Error updating S3 export path: {str(e)}")
+        raise
 

app/migrations/alembic_schema_models.py

@@ -23,6 +23,7 @@ class GenerationMetadataModel(Base):
     display_name = Column(Text)
     local_export_path = Column(Text)
     hf_export_path = Column(Text)
+    s3_export_path = Column(Text) 
     num_questions = Column(Float)
     total_count = Column(Float)
     topics = Column(Text)
@@ -66,6 +67,7 @@ class ExportMetadataModel(Base):
     display_name = Column(Text)
     local_export_path = Column(Text)
     hf_export_path = Column(Text)
+    s3_export_path = Column(Text) 
     job_id = Column(Text)
     job_name = Column(Text, unique=True)
     job_status = Column(Text)

app/models/request_models.py

@@ -45,10 +45,12 @@ class Example_eval(BaseModel):
     )
 
 
+# In app/models/request_models.py
 class S3Config(BaseModel):
     """S3 export configuration"""
     bucket: str
-    key: str
+    key: str = ""  # Make key optional with default empty string
+    create_if_not_exists: bool = True  # Flag to create bucket if it doesn't exist
 
 class HFConfig(BaseModel):
     """HF export configuration"""
@@ -59,41 +61,40 @@ class HFConfig(BaseModel):
     hf_commit_message: Optional[str] = "Hugging face export"  # Commit message
 
 class Export_synth(BaseModel):
-    # Export configuration
-    export_type: List[str] = Field(default_factory=lambda: ["huggingface"])  # Accept multiple export types (e.g., ["s3", "huggingface"])
-    file_path:str
-    display_name:Optional[str]= None
+    # Existing fields...
+    export_type: List[str] = Field(default_factory=lambda: ["huggingface"])
+    file_path: str
+    display_name: Optional[str] = None
     output_key: Optional[str] = 'Prompt'
     output_value: Optional[str] = 'Completion'
 
-
     # Hugging Face-specific fields
-    hf_config:HFConfig
+    hf_config: Optional[HFConfig] = None  # Make HF config optional
 
     # Optional s3 config
     s3_config: Optional[S3Config] = None
 
-    model_config = ConfigDict(protected_namespaces=(),
+    model_config = ConfigDict(
+        protected_namespaces=(),
         json_schema_extra={
             "example": {
-                        "export_type": [
-                                        "huggingface"
-                                        ],
-                    "file_path": "qa_pairs_claude_20241204_132411_test.json",
-                    "hf_config": {
-                        "hf_token": "your token",
-                        "hf_username": "your_username",
-                        "hf_repo_name": "file_name",
-                        "hf_commit_message": "dataset trial"
-                    }
-                
-                
+                "export_type": ["huggingface", "s3"],
+                "file_path": "qa_pairs_claude_20241204_132411_test.json",
+                "hf_config": {
+                    "hf_token": "your token",
+                    "hf_username": "your_username",
+                    "hf_repo_name": "file_name",
+                    "hf_commit_message": "dataset trial"
+                },
+                "s3_config": {
+                    "bucket": "my-dataset-bucket",
+                    "create_if_not_exists": True
+                }
             }
         }
     )
 
 
-
 class ModelParameters(BaseModel):
     """Low-level model parameters"""
     temperature: float = Field(default=0.0, ge=0.0, le=2.0, description="Controls randomness (0.0 to 1.0)")

app/run_export_job.py

@@ -8,24 +8,6 @@
 
     os.chdir("/home/cdsw/synthetic-data-studio")
 
-# def check_and_install_requirements():
-#     """Check and install requirements from requirements.txt"""
-#     # Get the current working directory instead of using __file__
-#     current_dir = os.getcwd()
-#     requirements_path = os.path.join(current_dir, 'requirements.txt')
-    
-#     if os.path.exists(requirements_path):
-#         try:
-#             print(f"Installing requirements from: {requirements_path}")
-#             subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-r', requirements_path])
-#         except subprocess.CalledProcessError as e:
-#             print(f"Error installing requirements: {e}")
-#             sys.exit(1)
-#     else:
-#         print("No requirements.txt found, continuing with existing packages")
-
-# # Run installation check at start
-# check_and_install_requirements()
 
 # Get the current notebook's directory
 notebook_dir = os.getcwd()

app/services/export_results.py

@@ -15,6 +15,7 @@
 from app.models.request_models import  Export_synth
 
 from app.core.database import DatabaseManager
+from app.services.s3_export import export_to_s3
 
 import logging
 from logging.handlers import RotatingFileHandler
@@ -101,38 +102,69 @@ def _create_dataset(self, records:List, output_key, output_value, file_path) ->
         return dataset
 
 
-    def export(self,request:Export_synth):
+    def export(self, request: Export_synth):
         try:
             export_paths = {}
             file_name = os.path.basename(request.file_path)
-            try:
-                with open(request.file_path, 'r') as f:
-                    output_data = json.load(f)
-            except FileNotFoundError:
-                raise HTTPException(status_code=404, detail=f"File not found: {request.file_path}")
-            except json.JSONDecodeError as e:
-                raise HTTPException(status_code=400, detail=f"Invalid JSON file: {str(e)}")
 
             for export_type in request.export_type:
-                if export_type == "s3" and request.s3_config:
-                    s3_client = boto3.client("s3")
-                    s3_client.put_object(
-                        Bucket=request.s3_config.bucket,
-                        Key=request.s3_config.key,
-                        Body=json.dumps(output_data, indent=2),
-                    )
-                    export_paths['s3']= f"s3://{request.s3_config.bucket}/{request.s3_config.key}"
-                    self.logger.info(f"Results saved to S3: {export_paths['s3']}")
-
+                # S3 Export
+                if export_type == "s3":
+                    if not request.s3_config:
+                        raise HTTPException(status_code=400, detail="S3 configuration required for S3 export")
+                    
+                    try:
+                        # Get bucket and key from request
+                        bucket_name = request.s3_config.bucket
+                        key = request.s3_config.key or file_name
+                        
+                        # Override with display_name if provided
+                        if request.display_name and not request.s3_config.key:
+                            key = f"{request.display_name}.json"
+                        
+                        
+                        
+                        
+                        create_bucket = getattr(request.s3_config, 'create_if_not_exists', True)
+                        
+                        s3_result = export_to_s3(
+                            file_path=request.file_path,
+                            bucket_name=bucket_name,
+                            key=key,
+                            create_bucket=create_bucket
+                        )
+                        
+                        s3_path = s3_result['s3']
+                        self.logger.info(f"Results saved to S3: {s3_path}")
+                        
+                        # Update database with S3 path
+                        self.db.update_s3_path(file_name, s3_path)
+                        self.logger.info(f"Generation Metadata updated for s3_path: {s3_path}")
+                        
+                        export_paths['s3'] = s3_path
+                        
+                    except Exception as e:
+                        self.logger.error(f"Error exporting to S3: {str(e)}", exc_info=True)
+                        raise APIError(f"S3 export failed: {str(e)}")
+                
+                # HuggingFace Export (existing code)
                 elif export_type == "huggingface" and request.hf_config:
+                    # We still need to read the file for HuggingFace export
+                    try:
+                        with open(request.file_path, 'r') as f:
+                            output_data = json.load(f)
+                    except FileNotFoundError:
+                        raise HTTPException(status_code=404, detail=f"File not found: {request.file_path}")
+                    except json.JSONDecodeError as e:
+                        raise HTTPException(status_code=400, detail=f"Invalid JSON file: {str(e)}")
+                    
                     self.logger.info(f"Creating HuggingFace dataset: {request.hf_config.hf_repo_name}")
 
                     # Set up HuggingFace authentication
                     HfFolder.save_token(request.hf_config.hf_token)
 
                     # Convert JSON to dataset
                     dataset = self._create_dataset(output_data, request.output_key, request.output_value, request.file_path)
-                    print(dataset)
 
                     # Push to HuggingFace Hub as a dataset
                     repo_id = f"{request.hf_config.hf_username}/{request.hf_config.hf_repo_name}"
@@ -146,8 +178,8 @@ def export(self,request:Export_synth):
                     self.logger.info(f"Dataset published to HuggingFace: {export_paths['huggingface']}")
                     self.db.update_hf_path(file_name, export_paths['huggingface'])
                     self.logger.info(f"Generation Metadata updated for hf_path: {export_paths['huggingface']}")
-                        
-                return export_paths
+            
+            return export_paths
 
         except Exception as e:
             self.logger.error(f"Error saving results: {str(e)}", exc_info=True)