FileNotFoundError: [Errno 2] No such file or directory: '/tmp/mds_data/train/index.json'

Hi guys i am using Mosaic ML package and trying to create data loader for training an LSTM model but encountering this error: FileNotFoundError: [Errno 2] No such file or directory: '/tmp/mds_data/train/index.json'
Please do share your feedback, if any one knows how to resolve this error?
from pyspark.sql import functions as F
from pyspark.ml.feature import MinMaxScaler, VectorAssembler, OneHotEncoder, StringIndexer
from pyspark.ml import Pipeline
from pyspark.sql.window import Window
from pyspark.ml.linalg import Vectors
from pyspark.sql.types import ArrayType, FloatType, IntegerType, FloatType, DoubleType, BinaryType,LongType, StringType
from pyspark.sql.functions import udf
import os
import shutil
import json
from pyspark.ml.functions import vector_to_array
import numpy as np
import pandas as pd
from streaming.base.converters import dataframe_to_mds
from streaming import StreamingDataset, StreamingDataLoader
import tempfile
from pyspark.sql.functions import flatten, col
from streaming.base.converters.dataframe_to_mds import dataframe_to_mds
from pyspark.dbutils import DBUtils
from pyspark.sql import SparkSession

class SparkDataPreprocessor:
def init(self, spark, train_file_path: str):
self.spark = spark
self.train_file_path = train_file_path
self.train_df = None
self.validation_df = None
self.test_df = None
self.scaler_model = None
self.encoder_model = None
self.num_features = 23
self.sequence_length = 50
self.num_classes = 3
self.remote_base = "/tmp/mds_data"
#self.remote_base = "dbfs:/mnt/deds_dv_data_products/wind_tum/Downloader"

def preprocess(self):
    self._configure_spark()
    self._load_and_prepare_data()
    self._split_data()
    self._normalize_data()
    self._generate_sequences()
    self._generate_labels()
    return self._create_streaming_dataloaders()
    #return self._collect_results()


def _configure_spark(self):
    self.spark.conf.set("spark.sql.shuffle.partitions", "400")
    self.spark.conf.set("spark.sql.adaptive.enabled", "true")
    #self.spark.conf.set("spark.memory.fraction", "0.8")
    #self.spark.conf.set("spark.memory.storageFraction", "0.3")
    #self.spark.conf.set("spark.executor.memory", "8g")
    #self.spark.conf.set("spark.driver.memory", "14g")
    self.spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
    self.spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "20000")


def _convert_to_mds(self, df, name):
    output_dir = f"/tmp/mds_data/{name}"                

    # Remove existing directory and create new one
    if os.path.exists(output_dir):
        print(f"Removing existing directory: {output_dir}")
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    

    # Check sequence structure before conversion
    sample = df.select("dummy_labels_seq").first()[0]
    print(len(sample), len(sample[0]))  # Should be (50, 3)

    #serialize_seq = udf(lambda arr: np.array(arr, dtype=np.float32).tobytes(),BinaryType())
    
    # Labels: 50x3 int64 array
    #serialize_labels = udf(lambda arr: np.array(arr, dtype=np.int64).tobytes(),BinaryType())
    #serialize_seq = udf(lambda arr: np.array(arr, dtype=np.float32).tolist(), ArrayType(ArrayType(FloatType())))


    #serialize_labels = udf(lambda arr: np.array(arr, dtype=np.int64).tolist(), StringType())
    #serialize_seq = udf(lambda arr: np.array(arr, dtype=np.float32).tobytes(), StringType())

    #serialize_labels = udf(lambda arr: json.dumps(arr), StringType())
    

    serialize_seq = udf(lambda arr: np.array(arr, dtype=np.float32).tobytes(), BinaryType())
    serialize_labels = udf(lambda arr: np.array(arr, dtype=np.int64).tobytes(), BinaryType())




    df = df.withColumn("sequence", serialize_seq("sequence"))
    df = df.withColumn("dummy_labels_seq", serialize_labels("dummy_labels_seq"))



    if df.rdd.isEmpty():
        raise ValueError(f"DataFrame '{name}' is empty - aborting MDS conversion")
    print(f"Number of rows in {name} DataFrame: {df.count()}")


    #if "sequence" not in df.columns:
    # raise ValueError("'sequence' column missing. Check _generate_sequences()")
    
    dataframe_to_mds(
        df.select("sequence", "dummy_labels_seq"),
        merge_index=0,
        mds_kwargs={
            'out': output_dir,
            'columns': {
                'sequence': 'bytes',
                'dummy_labels_seq': 'bytes',
                #'sequence': 'list:float32',
                #'dummy_labels_seq': 'list:int64'
            },
            'compression': 'zstd',
            'keep_local': True
        }
    )
     # Debug: Check if index.json exists after writing
    index_path = os.path.join(output_dir, "index.json")
    if os.path.exists(index_path):
        print(f"✅ index.json successfully created at {index_path}")
    else:
        print(f"❌ ERROR: index.json missing at {index_path}")
        
    return output_dir
    



"""
def _create_streaming_dataloaders(self):
    # Convert each dataset split to MDS format
    train_path = self._convert_to_mds(self.train_df, "train")       
    val_path = self._convert_to_mds(self.validation_df, "val")     
    test_path = self._convert_to_mds(self.test_df, "test")       

    #for path, name in [(train_path, "train"), (val_path, "val"), (test_path, "test")]:
    #    index_path = os.path.join(path, "index.json")
    #    if not os.path.exists(index_path):
    #        raise FileNotFoundError(f"index.json not found for {name} dataset at {index_path}")
    #    else:
    #        print(f"index.json found for {name} dataset at {index_path}")   
    
    # Return StreamingDataLoader instances
    return (
        StreamingDataLoader(StreamingDataset(local=train_path), batch_size=1024,predownload=2048),
        StreamingDataLoader(StreamingDataset(local=val_path), batch_size=1024,predownload=2048),
        StreamingDataLoader(StreamingDataset(local=test_path), batch_size=1024,predownload=2048)
    )

"""
def _create_streaming_dataloaders(self):
    def convert_and_create_dataset(df, name):
        # Convert DataFrame to MDS and get the directory path
        #df.select("sequence", "dummy_labels_seq").show(5, truncate=False)

        output_dir = self._convert_to_mds(df, name)

        # Ensure index.json exists before proceeding
        index_path = os.path.join(output_dir, "index.json")
        #if not os.path.exists(index_path):
        #    raise FileNotFoundError(f"index.json not found for {name} dataset at {index_path}")
        
        # Create StreamingDataset
        return StreamingDataset(remote=output_dir, local=output_dir, split=None)  

    # Create datasets
    train_dataset = convert_and_create_dataset(self.train_df, "train")
    val_dataset = convert_and_create_dataset(self.validation_df, "val")
    test_dataset = convert_and_create_dataset(self.test_df, "test")

    # Create dataloaders
    return (
        StreamingDataLoader(train_dataset, batch_size=32, predownload= 256),#batch 1024,  predownload=2048
        StreamingDataLoader(val_dataset, batch_size=32, predownload=256),
        StreamingDataLoader(test_dataset, batch_size=32, predownload=256)
    )



def _load_and_prepare_data(self):
    """Load data with optimized partitioning"""
    self.train_df = (self.spark.read.table(self.train_file_path)
                    .repartition(400, "turbine_status_id")
                    .sortWithinPartitions("turbine_status_id", "timestamp")
                    .cache())
    #self.train_df.count()  # Force caching
    return self.train_df

def _split_data(self):
    window = Window.partitionBy("class_label") \
             .orderBy("turbine_status_id", "timestamp")

    # Add stratified-temporal rank to original data
    ranked_df = self.train_df.withColumn("rank", F.percent_rank().over(window))       
    # Split the data
    self.train_df = ranked_df.filter(F.col("rank") <= 0.7).drop("rank").cache()
    self.validation_df = ranked_df.filter((F.col("rank") > 0.7) & (F.col("rank") <= 0.85)).drop("rank").cache()
    self.test_df = ranked_df.filter(F.col("rank") > 0.85).drop("rank").cache()

    # Validate splits
    for df, name in [(self.train_df, "Training"), 
                    (self.validation_df, "Validation"),
                    (self.test_df, "Test")]:
        if df.isEmpty():
            raise ValueError(f"{name} DataFrame is empty after splitting")
        df.count() 

def _normalize_data(self):
    exclude_cols = {"turbine_status_id", "timestamp", "class_label"}
    numeric_cols = [
        'wind_speed', 'wind_direction', 'nacelle_position', 'power',
        'rotor_speed', 'blade_angle_pitch_position', 'ambient_temperature',
        'potential_power_default_pc', 'potential_power_learned_pc',
        'potential_power_mpc', 'time_based_system_avail',
        'production_based_system_avail', 'data_availability',
        'temperature_ambient_nacelle_inside', 'temperature_ambient_nacelle_outside',
        'temperature_generator_rotor_1', 'temperature_generator_rotor_2',
        'temperature_generator_stator_1', 'temperature_generator_stator_2',
        'temperature_rotor_bearing_front', 'temperature_rotor_bearing_rear',
        'temperature_spinner', 'temperature_transformer_default']
    
    # Validate numeric columns count
    print(f"Normalizing {len(numeric_cols)} features: {numeric_cols}")
    
    assembler = VectorAssembler(inputCols=numeric_cols, outputCol="features_vec")
    scaler = MinMaxScaler(inputCol="features_vec", outputCol="scaled_features_vec")
    pipeline = Pipeline(stages=[assembler, scaler])
    
    self.scaler_model = pipeline.fit(self.train_df)
    
    for df_attr in ['train_df', 'validation_df', 'test_df']:
        df = getattr(self, df_attr)
        df = self.scaler_model.transform(df)
        df = df.drop("features_vec", *numeric_cols)  # Remove original numeric columns
        # Select only required columns
        #df = df.select(*exclude_cols,"scaled_features_vec") # i think only "scaled_features_vec" will be considered
        df = df.select("scaled_features_vec",*exclude_cols)
        
        setattr(self, df_attr, df.cache())




def _generate_sequences(self):
    sequence_length = 50
    #window = Window.partitionBy("turbine_status_id") \
    #             .orderBy("timestamp") \
    #             .rowsBetween(-sequence_length + 1, 0)

    window = Window.rowsBetween(-sequence_length + 1, 0)

    # Get actual feature count from the vector
    sample_vec = self.train_df.select("scaled_features_vec").first()[0]
    self.num_features = len(sample_vec)  # Should be 23
    print(f"Confirmed features per timestep: {self.num_features}")
    
    for df_attr in ['train_df', 'validation_df', 'test_df']:
        df = getattr(self, df_attr)
        df = df.withColumn("sequence", 
            F.collect_list(
                vector_to_array("scaled_features_vec").cast("array<float>")
            ).over(window)
        ).filter(F.size("sequence") == sequence_length)
        
        # Validate sequences
        sample = df.select("sequence").first()[0]
        assert len(sample) == 50 and len(sample[0]) == self.num_features, \
            f"Invalid sequence shape: {len(sample)}x{len(sample[0])} (expected 50x{self.num_features})"

        print(len(sample), len(sample[0]))  # Should be (50, 3)
        print(f"data type of sequence: {type(sample[0])}")
        setattr(self, df_attr, df.drop("scaled_features_vec").cache())

    

def _generate_labels(self):
    """Efficient one-hot encoding with model reuse"""
    #self.train_df =self._load_and_prepare_data()
    indexer = StringIndexer(inputCol="class_label", outputCol="label_index")
    encoder = OneHotEncoder(inputCol="label_index", outputCol="dummy_labels", dropLast=False)
    #encoder = OneHotEncoder(inputCol="class_label", outputCol="dummy_labels", dropLast=False)
    pipeline = Pipeline(stages=[indexer, encoder])
    #pipeline = Pipeline(stages=[encoder])
    self.encoder_model = pipeline.fit(self.train_df)
    
    #window = Window.partitionBy("turbine_status_id") \
    #            .orderBy("timestamp") \
    #            .rowsBetween(-self.sequence_length + 1, 0)
    
    window = Window.rowsBetween(-self.sequence_length + 1, 0)

    for df_attr in ['train_df', 'validation_df', 'test_df']:
        df = getattr(self, df_attr)
        df = self.encoder_model.transform(df)
        
        df = df.withColumn("dummy_labels_seq", 
        F.collect_list(vector_to_array("dummy_labels")).over(window)).filter(F.size("dummy_labels_seq") == self.sequence_length)
    
        # Keep necessary columns
        df = df.select("sequence", "dummy_labels_seq", "turbine_status_id", "timestamp")
        # Check sequence structure before conversion
        sample = df.select("dummy_labels_seq").first()[0]
        print(len(sample), len(sample[0]))  # Should be (50, 3)
        print(f"data type of dummy labels: {type(sample[0])}")
        setattr(self, df_attr, df)




def _collect_results(self):
    def process_df(df):
        return (
            np.array(df.select("sequence").collect()),
            np.array(df.select(vector_to_array("dummy_labels_vec")).collect())
        )
    
    train_seq, train_labels = process_df(self.train_df)
    val_seq, val_labels = process_df(self.validation_df)
    test_seq, test_labels = process_df(self.test_df)

    return train_seq, train_labels, val_seq, val_labels, test_seq, test_labels

Start in _convert_to_mds and _create_streaming_dataloaders, then trace the call to dataframe_to_mds and the StreamingDataset construction for the train split. Reproduce the failure with the provided preprocessing flow and verify that the generated train output contains the expected index.json and can be opened by StreamingDataset.