Untitled

    import os
    
    checkpoint_dir = "/spark-local-data/checkpoints"
    
    print("\n" + "=" * 80)
    print("[CHECKPOINT] Setting up checkpoint directory...")
    print("=" * 80)
    print(f"[CHECKPOINT] Target directory: {checkpoint_dir}")
    
    # Step 1: Create directory if not exists
    try:
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir, mode=0o755, exist_ok=True)
            print(f"[CHECKPOINT] ✓ Created directory: {checkpoint_dir}")
        else:
            print(f"[CHECKPOINT] ✓ Directory already exists: {checkpoint_dir}")
            
        # Verify directory is writable
        test_file = os.path.join(checkpoint_dir, ".test_write")
        with open(test_file, 'w') as f:
            f.write("test")
        os.remove(test_file)
        print(f"[CHECKPOINT] ✓ Directory is writable")
        
    except Exception as e:
        print(f"[CHECKPOINT] ⚠ Warning: Directory creation/test failed: {e}")
        print(f"[CHECKPOINT] Will attempt to continue...")
    
    # Step 2: Set checkpoint directory in SparkContext (MANDATORY!)
    try:
        spark.sparkContext.setCheckpointDir(checkpoint_dir)
        print(f"[CHECKPOINT] ✓ Checkpoint directory set in SparkContext")
        
    except Exception as e:
        print(f"[CHECKPOINT] ✗ CRITICAL ERROR: Failed to set checkpoint directory!")
        print(f"[CHECKPOINT] Error: {e}")
        print(f"[CHECKPOINT] Job cannot continue without checkpoint support")
        raise RuntimeError(f"Checkpoint setup failed: {e}")
    
    # Step 3: Verify checkpoint directory is properly set
    try:
        current_checkpoint_dir = spark.sparkContext.getCheckpointDir()
        if current_checkpoint_dir:
            # getCheckpointDir returns Optional, need to extract value
            actual_dir = current_checkpoint_dir.get() if hasattr(current_checkpoint_dir, 'get') else str(current_checkpoint_dir)
            print(f"[CHECKPOINT] ✓ Verification successful")
            print(f"[CHECKPOINT] ✓ Active checkpoint directory: {actual_dir}")
        else:
            print(f"[CHECKPOINT] ✗ WARNING: Checkpoint directory not confirmed")
            print(f"[CHECKPOINT] Will attempt to continue...")
            
    except Exception as e:
        print(f"[CHECKPOINT] ⚠ Could not verify checkpoint directory: {e}")
        print(f"[CHECKPOINT] But setCheckpointDir was called, should be OK")
    
    print("=" * 80)
    print("[CHECKPOINT] ✓ Checkpoint setup complete")
    print("=" * 80 + "\n")