Hi @anirbanmishra 

I see - you're logging metrics during training (which works fine) but encountering the file descriptor/Py4J gateway errors specifically when logging artifacts at the end after 350 epochs. This is actually a more complex issue because by that point, you likely have resource exhaustion from the long-running training process. Here are targeted solutions for your specific scenario:

Solution 1: Reset Spark Context Before Artifact Logging

import mlflow
from pyspark.sql import SparkSession

# After training completes but before artifact logging
def reset_spark_context():
try:
# Get current Spark context
current_spark = SparkSession.getActiveSession()
if current_spark:
current_spark.stop()

# Wait a moment for cleanup
import time
time.sleep(5)

# Create fresh Spark session
new_spark = SparkSession.builder \
.appName("ArtifactLogging") \
.config("spark.sql.adaptive.enabled", "true") \
.config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
.getOrCreate()

return new_spark
except Exception as e:
print(f"Error resetting Spark context: {e}")
return None

# Your training workflow
with mlflow.start_run() as run:
# Training loop with metric logging
for epoch in range(350):
# ... training code ...
mlflow.log_metrics({"loss": loss, "accuracy": acc}, step=epoch)

print("Training completed, preparing to log artifacts...")

# Reset Spark context before artifact logging
spark = reset_spark_context()

# Now log artifacts
print("Logging artifacts...")
mlflow.log_artifacts("path/to/artifacts", "artifacts")

Solution 2: Chunked Artifact Logging with Delays

import time
import os
import mlflow

def log_artifacts_in_chunks(artifact_dir, chunk_size=5, delay_seconds=2):
"""Log artifacts in small chunks with delays to prevent resource exhaustion"""

artifact_files = []
for root, dirs, files in os.walk(artifact_dir):
for file in files:
artifact_files.append(os.path.join(root, file))

print(f"Total artifacts to log: {len(artifact_files)}")

# Process in chunks
for i in range(0, len(artifact_files), chunk_size):
chunk = artifact_files[i:i+chunk_size]

print(f"Logging chunk {i//chunk_size + 1}/{(len(artifact_files)-1)//chunk_size + 1}")

for artifact_path in chunk:
try:
mlflow.log_artifact(artifact_path)
print(f"Logged: {os.path.basename(artifact_path)}")
except Exception as e:
print(f"Failed to log {artifact_path}: {e}")

# Delay between chunks to let resources recover
time.sleep(delay_seconds)

# Usage after training
with mlflow.start_run() as run:
# ... training code ...

print("Training completed, logging artifacts in chunks...")
log_artifacts_in_chunks("path/to/artifacts", chunk_size=3, delay_seconds=3)

This approach should resolve the file descriptor and Py4J gateway issues you're experiencing while ensuring all your artifacts are properly logged for downstream analysis.