hunk.code_window: import ray
   from ray._private.test_utils import wait_for_condition
   from ray.cluster_utils import Cluster, cluster_not_supported
   from ray.tests.test_object_spilling import assert_no_thrashing, is_dir_empty
   @pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
   def test_multiple_directories(tmp_path, shutdown_only):
       num_dirs = 3
  
hunk.labels: keep keep keep keep add keep keep keep keep
hunk.after_edit: # Note: Disk write speed can be as low as 6 MiB/s in AWS Mac instances, so we have to
   # increase the timeout.
   pytestmark = [pytest.mark.timeout(900 if platform.system() == "Darwin" else 180)]
  
hunk.file_path: python/ray/tests/test_object_spilling_3.py
hunk.type: add
file: import unittest
  
  import numpy as np
  
  import ray
  from ray.rllib.policy.rnn_sequencing import (
      pad_batch_to_sequences_of_same_size,
      add_time_dimension,
      chop_into_sequences,
  )
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.test_utils import check
  
  tf1, tf, tfv = try_import_tf()
  torch, nn = try_import_torch()
  
  
  class TestRNNSequencing(unittest.TestCase):
      @classmethod
      def setUpClass(cls) -> None:
          ray.init()
  
      @classmethod
      def tearDownClass(cls) -> None:
          ray.shutdown()
  
      def test_chop_into_sequences_long_seq(self):
          """Test pad_batch where episodes are longer than max_seq_len. The long
          seq should be split into two smaller seqs that are less than max_seq_len"""
          max_seq_len = 2
          # Input seq lens, corresponding to ep_ids, unroll_ids, etc.
          seq_lens = [2, 3, 1]  # noqa: F841
          ep_ids = [0, 0, 1, 1, 1, 2]
          unroll_ids = [2, 2, 3, 3, 3, 4]
          feats = [[1, 1, 2, 2, 2, 3]]
          # Input states, ie states[3] is the input state at
          # t = 3 and the output state at t = 2
          states = [[1, 2, 3, 4, 5, 6]]
          agent = [0, 0, 0, 0, 0, 0]
          f_pad, s_init, s_lens = chop_into_sequences(
              feature_columns=feats,
              state_columns=states,
              max_seq_len=max_seq_len,
              episode_ids=ep_ids,
              unroll_ids=unroll_ids,
              agent_indices=agent,
              dynamic_max=False,
              padding="zero",
          )
          expected_f_pad = [[1, 1, 2, 2, 2, 0, 3, 0]]
          expected_seq_lens = [2, 2, 1, 1]
          expected_states = [[1, 3, 5, 6]]
          check(f_pad, expected_f_pad)
          check(s_lens, expected_seq_lens)
          check(s_init, expected_states)
  
          # Try again with dynamic max
          f_pad, s_init, s_lens = chop_into_sequences(
              feature_columns=feats,
              state_columns=states,
              max_seq_len=max_seq_len,
              episode_ids=ep_ids,
              unroll_ids=unroll_ids,
              agent_indices=agent,
              dynamic_max=True,
          )
          check(f_pad, expected_f_pad)
          check(s_lens, expected_seq_lens)
          check(s_init, expected_states)
  
          # With padding="last"
          f_pad, s_init, s_lens = chop_into_sequences(
              feature_columns=feats,
              state_columns=states,
              max_seq_len=max_seq_len,
              episode_ids=ep_ids,
              unroll_ids=unroll_ids,
              agent_indices=agent,
              dynamic_max=False,
              padding="last",
          )
          expected_f_pad = [[1, 1, 2, 2, 2, 2, 3, 3]]
          check(f_pad, expected_f_pad)
          check(s_lens, expected_seq_lens)
          check(s_init, expected_states)
  
      def test_pad_batch_dynamic_max(self):
          """Test pad_batch_to_sequences_of_same_size when dynamic_max = True"""
          view_requirements = {
              "state_in_0": ViewRequirement(
                  "state_out_0",
                  shift=[-1],
                  used_for_training=False,
                  used_for_compute_actions=True,
                  batch_repeat_value=1,
              )
          }
          max_seq_len = 20
          num_seqs = np.random.randint(1, 20)
          seq_lens = np.random.randint(1, max_seq_len, size=(num_seqs))
          max_len = np.max(seq_lens)
          sum_seq_lens = np.sum(seq_lens)
  
          s1 = SampleBatch(
              {
                  "a": np.arange(sum_seq_lens),
                  "b": np.arange(sum_seq_lens),
                  "seq_lens": seq_lens,
                  "state_in_0": [[0]] * num_seqs,
              },
              _max_seq_len=max_seq_len,
          )
  
          pad_batch_to_sequences_of_same_size(
              s1,
              max_seq_len=max_seq_len,
              feature_keys=["a", "b"],
              view_requirements=view_requirements,
          )
          check(s1.max_seq_len, max_len)
          check(s1["a"].shape[0], max_len * num_seqs)
          check(s1["b"].shape[0], max_len * num_seqs)
  
      def test_pad_batch_fixed_max(self):
          """Test pad_batch_to_sequences_of_same_size when dynamic_max = False"""
          view_requirements = {
              "state_in_0": ViewRequirement(
                  "state_out_0",
                  shift="-3:-1",
                  used_for_training=False,
                  used_for_compute_actions=True,
                  batch_repeat_value=1,
              )
          }
          max_seq_len = 20
          num_seqs = np.random.randint(1, 20)
          seq_lens = np.random.randint(1, max_seq_len, size=(num_seqs))
          sum_seq_lens = np.sum(seq_lens)
          s1 = SampleBatch(
              {
                  "a": np.arange(sum_seq_lens),
                  "b": np.arange(sum_seq_lens),
                  "seq_lens": seq_lens,
                  "state_in_0": [[0]] * num_seqs,
              },
              _max_seq_len=max_seq_len,
          )
  
          pad_batch_to_sequences_of_same_size(
              s1,
              max_seq_len=max_seq_len,
              feature_keys=["a", "b"],
              view_requirements=view_requirements,
          )
          check(s1.max_seq_len, max_seq_len)
          check(s1["a"].shape[0], max_seq_len * num_seqs)
          check(s1["b"].shape[0], max_seq_len * num_seqs)
  
      def test_add_time_dimension(self):
          """Test add_time_dimension gives sequential data along the time dimension"""
  
          B, T, F = np.random.choice(
              np.asarray(list(range(8, 32)), dtype=np.int32),  # use int32 for seq_lens
              size=3,
              replace=False,
          )
  
          inputs_numpy = np.repeat(
              np.arange(B * T)[:, np.newaxis], repeats=F, axis=-1
          ).astype(np.int32)
          check(inputs_numpy.shape, (B * T, F))
          seq_lens = np.random.randint(1, T + 1, (B,))
  
          time_shift_diff_batch_major = np.ones(shape=(B, T - 1, F), dtype=np.int32)
          time_shift_diff_time_major = np.ones(shape=(T - 1, B, F), dtype=np.int32)
  
          if tf is not None:
              # Test tensorflow batch-major
              padded_inputs = tf.constant(inputs_numpy)
              batch_major_outputs = add_time_dimension(
                  padded_inputs, seq_lens=seq_lens, framework="tf", time_major=False
              )
              check(batch_major_outputs.shape.as_list(), [B, T, F])
              time_shift_diff = batch_major_outputs[:, 1:] - batch_major_outputs[:, :-1]
              check(time_shift_diff, time_shift_diff_batch_major)
  
          if torch is not None:
              # Test torch batch-major
              padded_inputs = torch.from_numpy(inputs_numpy)
              batch_major_outputs = add_time_dimension(
                  padded_inputs, seq_lens=seq_lens, framework="torch", time_major=False
              )
              check(batch_major_outputs.shape, (B, T, F))
              time_shift_diff = batch_major_outputs[:, 1:] - batch_major_outputs[:, :-1]
              check(time_shift_diff, time_shift_diff_batch_major)
  
              # Test torch time-major
              padded_inputs = torch.from_numpy(inputs_numpy)
              time_major_outputs = add_time_dimension(
                  padded_inputs, seq_lens=seq_lens, framework="torch", time_major=True
              )
              check(time_major_outputs.shape, (T, B, F))
              time_shift_diff = time_major_outputs[1:] - time_major_outputs[:-1]
              check(time_shift_diff, time_shift_diff_time_major)
  
  
  if __name__ == "__main__":
      import pytest
      import sys
  
      sys.exit(pytest.main(["-v", __file__]))
  