Fix multiprocessing segfault

Author: coolkp

.github/workflows/UnitTests.yml

@@ -53,8 +53,8 @@ jobs:
       run: |
         ruff check .
     - name: PyTest
-      run: | 
-        HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py --deselect=src/maxdiffusion/tests/input_pipeline_interface_test.py -x
+      run: | #--deselect=src/maxdiffusion/tests/input_pipeline_interface_test.py
+        HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py  -x
 #  add_pull_ready:
 #    if: github.ref != 'refs/heads/main'
 #    permissions:

src/maxdiffusion/input_pipeline/_tfds_data_processing.py

@@ -22,7 +22,7 @@
 from maxdiffusion import multihost_dataloading, max_logging
 
 AUTOTUNE = tf.data.AUTOTUNE
-
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 def load_as_tf_dataset(dataset, global_batch_size, shuffle, dataloading_host_count):
   dataset = dataset.with_format("tensorflow")[:]

src/maxdiffusion/input_pipeline/input_pipeline_interface.py

@@ -40,7 +40,7 @@
 from PIL import Image
 
 AUTOTUNE = tf.data.experimental.AUTOTUNE
-
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 def make_data_iterator(
     config,

src/maxdiffusion/tests/input_pipeline_interface_test.py

@@ -23,7 +23,6 @@
 from absl.testing import absltest
 
 import numpy as np
-import pytest
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp
 import jax
@@ -70,7 +69,6 @@ class InputPipelineInterface(unittest.TestCase):
   def setUp(self):
     InputPipelineInterface.dummy_data = {}
 
-  @pytest.mark.skip(reason="Debug segfault")
   def test_make_dreambooth_train_iterator(self):
 
     instance_class_gcs_dir = "gs://maxdiffusion-github-runner-test-assets/datasets/dreambooth/instance_class"