Merge branch 'master' into lkchen-ray_data_llm

Signed-off-by: Linkun Chen <github@lkchen.net>

Author: lk-chen

.buildkite/others.rayci.yml

@@ -19,6 +19,14 @@ steps:
     job_env: oss-ci-base_test-py3.11
     depends_on: oss-ci-base_test-multipy
 
+  - label: ":tapioca: build: uv pip compile LLM dependencies"
+    key: uv_pip_compile_llm_dependencies
+    instance_type: small
+    command: ./ci/test_compile_llm_requirements.sh
+    soft_fail: true
+    job_env: oss-ci-base_test-py3.11
+    depends_on: oss-ci-base_test-multipy
+
   # docs
   - name: doctestbuild
     wanda: ci/docker/doctest.build.wanda.yaml

ci/compile_llm_requirements.sh

@@ -3,21 +3,31 @@
 set -euo pipefail
 
 PYTHON_CODE="$(python -c "import sys; v=sys.version_info; print(f'py{v.major}{v.minor}')")"
+if [[ "${PYTHON_CODE}" != "py311" ]]; then
+	echo "--- Python version is not 3.11"
+	echo "--- Current Python version: ${PYTHON_CODE}"
+	exit 1
+fi
 
 for CUDA_CODE in cpu cu121 cu124 ; do
 	PYTHON_CUDA_CODE="${PYTHON_CODE}_${CUDA_CODE}"
 
 	echo "--- Compile dependencies for ${PYTHON_CODE}_${CUDA_CODE}"
 
-	PIP_COMPILE=(
-		pip-compile -v --generate-hashes --strip-extras
+	UV_PIP_COMPILE=(
+		uv pip compile -v --generate-hashes --strip-extras
 		--unsafe-package ray
 		# The version we use on python 3.9 is not installable on python 3.11
 		--unsafe-package grpcio-tools
 		# setuptools should not be pinned.
 		--unsafe-package setuptools
+		--index-url "https://pypi.org/simple"
 		--extra-index-url "https://download.pytorch.org/whl/${CUDA_CODE}"
-		--find-links "https://data.pyg.org/whl/torch-2.3.0+${CUDA_CODE}.html"
+		--find-links "https://data.pyg.org/whl/torch-2.5.1+${CUDA_CODE}.html"
+		--index-strategy unsafe-best-match
+		--no-strip-markers
+		--emit-index-url
+		--emit-find-links
 	)
 
 	mkdir -p /tmp/ray-deps
@@ -32,7 +42,7 @@ for CUDA_CODE in cpu cu121 cu124 ; do
 	#
 	# Needs to use the exact torch version.
 	echo "--- Compile ray base test dependencies"
-	"${PIP_COMPILE[@]}" \
+	"${UV_PIP_COMPILE[@]}" \
 		-c "/tmp/ray-deps/requirements_compiled.txt" \
 		"python/requirements.txt" \
 		"python/requirements/cloud-requirements.txt" \
@@ -41,7 +51,7 @@ for CUDA_CODE in cpu cu121 cu124 ; do
 
 	# Second, expand it into LLM test dependencies
 	echo "--- Compile LLM test dependencies"
-	"${PIP_COMPILE[@]}" \
+	"${UV_PIP_COMPILE[@]}" \
 		-c "python/requirements_compiled_ray_test_${PYTHON_CUDA_CODE}.txt" \
 		"python/requirements.txt" \
 		"python/requirements/cloud-requirements.txt" \
@@ -53,7 +63,7 @@ for CUDA_CODE in cpu cu121 cu124 ; do
 	# Third, extract the ray base dependencies from ray base test dependencies.
 	# TODO(aslonnie): This should be used for installing ray in the container images.
 	echo "--- Compile ray base test dependencies"
-	"${PIP_COMPILE[@]}" \
+	"${UV_PIP_COMPILE[@]}" \
 		-c "python/requirements_compiled_ray_test_${PYTHON_CUDA_CODE}.txt" \
 		"python/requirements.txt" \
 		-o "python/requirements_compiled_ray_${PYTHON_CUDA_CODE}.txt"
@@ -62,10 +72,11 @@ for CUDA_CODE in cpu cu121 cu124 ; do
 	# which is also an expansion of the ray base dependencies.
 	# TODO(aslonnie): This should be used for installing ray[llm] in the container images.
 	echo "--- Compile LLM dependencies"
-	"${PIP_COMPILE[@]}" \
+	"${UV_PIP_COMPILE[@]}" \
 		-c "python/requirements_compiled_rayllm_test_${PYTHON_CUDA_CODE}.txt" \
 		"python/requirements.txt" \
 		"python/requirements/llm/llm-requirements.txt" \
 		-o "python/requirements_compiled_rayllm_${PYTHON_CUDA_CODE}.txt"
-
 done
+
+echo "--- Done"

ci/pipeline/determine_tests_to_run.py

@@ -11,6 +11,7 @@
 
 _ALL_TAGS = set(
     """
+    always
     lint python cpp core_cpp java workflow accelerated_dag dashboard
     data serve ml tune train llm rllib rllib_gpu rllib_directly
     linux_wheels macos_wheels docker doc python_dependencies tools
@@ -189,6 +190,7 @@ def match_tags(self, changed_file: str) -> Tuple[Set[str], bool]:
 
     tags: Set[str] = set()
 
+    tags.add("always")
     tags.add("lint")
 
     def _emit(line: str):

ci/pipeline/test_conditional_testing.py

@@ -153,7 +153,7 @@ def __init__(self, file: str, tags: Set[str]):
             )
             tags = output.split()
 
-            want = test_case.tags
+            want = set(list(test_case.tags) + ["always"])
             assert want == set(tags), f"file {test_case.file}, want {want}, got {tags}"
 
 

ci/test_compile_llm_requirements.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+set -e
+
+# Install uv and set up Python
+pip install uv
+uv python install 3.11
+uv python pin 3.11
+
+# Create a temporary directory for backup files and setup cleanup trap
+TEMP_DIR=$(mktemp -d)
+cleanup() {
+    echo "Cleaning up temporary directory: $TEMP_DIR"
+    rm -rf "$TEMP_DIR"
+}
+trap cleanup EXIT
+
+echo "Created temporary directory: $TEMP_DIR"
+
+# Create backup copies of req files to reference to
+cp ./python/requirements_compiled_rayllm_py311_cpu.txt "$TEMP_DIR/requirements_compiled_rayllm_py311_cpu_backup.txt"
+cp ./python/requirements_compiled_rayllm_py311_cu121.txt "$TEMP_DIR/requirements_compiled_rayllm_py311_cu121_backup.txt"
+cp ./python/requirements_compiled_rayllm_py311_cu124.txt "$TEMP_DIR/requirements_compiled_rayllm_py311_cu124_backup.txt"
+
+./ci/compile_llm_requirements.sh
+
+# Copy files to artifact mount on Buildkite
+cp ./python/requirements_compiled_rayllm_py311_cpu.txt /artifact-mount/
+cp ./python/requirements_compiled_rayllm_py311_cu121.txt /artifact-mount/
+cp ./python/requirements_compiled_rayllm_py311_cu124.txt /artifact-mount/
+
+# Check all files and print if files are not up to date
+FAILED=0
+for VARIANT in cpu cu121 cu124; do
+    diff --color -u ./python/requirements_compiled_rayllm_py311_${VARIANT}.txt "$TEMP_DIR/requirements_compiled_rayllm_py311_${VARIANT}_backup.txt" || {
+        echo "requirements_compiled_rayllm_py311_${VARIANT}.txt is not up to date. Please download it from Artifacts tab and git push the changes."
+        FAILED=1
+    }
+done
+if [[ $FAILED -eq 1 ]]; then
+    exit 1
+fi

python/ray/_private/state_api_test_utils.py

@@ -29,6 +29,7 @@
     PredicateType,
     SupportedFilterType,
 )
+import ray._private.test_utils as test_utils
 
 
 @dataclass
@@ -113,6 +114,16 @@ def invoke_state_api(
     return res
 
 
+def invoke_state_api_n(*args, **kwargs):
+    def verify():
+        NUM_API_CALL_SAMPLES = 10
+        for _ in range(NUM_API_CALL_SAMPLES):
+            invoke_state_api(*args, **kwargs)
+        return True
+
+    test_utils.wait_for_condition(verify, retry_interval_ms=2000, timeout=30)
+
+
 def aggregate_perf_results(state_stats: StateAPIStats = GLOBAL_STATE_STATS):
     """Aggregate stats of state API calls
 

python/ray/data/_internal/datasource/databricks_uc_datasource.py

@@ -37,7 +37,10 @@ def __init__(
         self.schema = schema
         self.query = query
 
-        url_base = f"https://{self.host}/api/2.0/sql/statements/"
+        if not host.startswith(("http://", "https://")):
+            self.host = f"https://{host}"
+
+        url_base = f"{self.host}/api/2.0/sql/statements/"
 
         payload = json.dumps(
             {

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -376,7 +376,7 @@ def _generate_transform_fn_for_map_batches(
 ) -> MapTransformCallable[DataBatch, DataBatch]:
     if inspect.iscoroutinefunction(fn):
         # UDF is a callable class with async generator `__call__` method.
-        transform_fn = _generate_transform_fn_for_async_map_batches(fn)
+        transform_fn = _generate_transform_fn_for_async_map(fn, _validate_batch_output)
 
     else:
 
@@ -423,64 +423,66 @@ def transform_fn(
     return transform_fn
 
 
-def _generate_transform_fn_for_async_map_batches(
+def _generate_transform_fn_for_async_map(
     fn: UserDefinedFunction,
-) -> MapTransformCallable[DataBatch, DataBatch]:
-    def transform_fn(
-        input_iterable: Iterable[DataBatch], _: TaskContext
-    ) -> Iterable[DataBatch]:
+    validate_fn,
+) -> MapTransformCallable:
+    # Generates a transform function for asynchronous mapping of items (either batches or rows)
+    # using a user-defined function (UDF). This consolidated function handles both asynchronous
+    # batch processing and asynchronous flat mapping (e.g., rows) based on the provided UDF.
+    def transform_fn(input_iterable: Iterable, _: TaskContext) -> Iterable:
         # Use a queue to store outputs from async generator calls.
-        # We will put output batches into this queue from async
+        # We will put output items into this queue from async
         # generators, and in the main event loop, yield them from
         # the queue as they become available.
-        output_batch_queue = queue.Queue()
+        output_item_queue = queue.Queue()
         # Sentinel object to signal the end of the async generator.
         sentinel = object()
 
-        async def process_batch(batch: DataBatch):
+        async def process_item(item):
             try:
-                output_batch_iterator = await fn(batch)
+                output_item_iterator = await fn(item)
                 # As soon as results become available from the async generator,
                 # put them into the result queue so they can be yielded.
-                async for output_batch in output_batch_iterator:
-                    output_batch_queue.put(output_batch)
+                async for output_item in output_item_iterator:
+                    output_item_queue.put(output_item)
             except Exception as e:
-                output_batch_queue.put(
+                output_item_queue.put(
                     e
                 )  # Put the exception into the queue to signal an error
 
-        async def process_all_batches():
+        async def process_all_items():
             try:
                 loop = ray.data._map_actor_context.udf_map_asyncio_loop
-                tasks = [loop.create_task(process_batch(x)) for x in input_iterable]
+                tasks = [loop.create_task(process_item(x)) for x in input_iterable]
 
                 ctx = ray.data.DataContext.get_current()
                 if ctx.execution_options.preserve_order:
                     for task in tasks:
-                        await task()
+                        await task
                 else:
                     for task in asyncio.as_completed(tasks):
                         await task
             finally:
-                output_batch_queue.put(sentinel)
+                output_item_queue.put(sentinel)
 
-        # Use the existing event loop to create and run Tasks to process each batch
+        # Use the existing event loop to create and run Tasks to process each item
         loop = ray.data._map_actor_context.udf_map_asyncio_loop
-        asyncio.run_coroutine_threadsafe(process_all_batches(), loop)
+        asyncio.run_coroutine_threadsafe(process_all_items(), loop)
 
         # Yield results as they become available.
         while True:
-            # Here, `out_batch` is a one-row output batch
+            # Here, `out_item` is a one-row output item
             # from the async generator, corresponding to a
-            # single row from the input batch.
-            out_batch = output_batch_queue.get()
-            if out_batch is sentinel:
+            # single row from the input item.
+            out_item = output_item_queue.get()
+            if out_item is sentinel:
                 # Break out of the loop when the sentinel is received.
                 break
-            if isinstance(out_batch, Exception):
-                raise out_batch
-            _validate_batch_output(out_batch)
-            yield out_batch
+            if isinstance(out_item, Exception):
+                raise out_item
+            validate_fn(out_item)
+            yield out_item
 
     return transform_fn
 
@@ -511,11 +513,17 @@ def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
 def _generate_transform_fn_for_flat_map(
     fn: UserDefinedFunction,
 ) -> MapTransformCallable[Row, Row]:
-    def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
-        for row in rows:
-            for out_row in fn(row):
-                _validate_row_output(out_row)
-                yield out_row
+    if inspect.iscoroutinefunction(fn):
+        # UDF is a callable class with async generator `__call__` method.
+        transform_fn = _generate_transform_fn_for_async_map(fn, _validate_row_output)
+
+    else:
+
+        def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+            for row in rows:
+                for out_row in fn(row):
+                    _validate_row_output(out_row)
+                    yield out_row
 
     return transform_fn
 

python/ray/data/dataset.py

@@ -1373,12 +1373,6 @@ def repartition(
         """Repartition the :class:`Dataset` into exactly this number of
         :ref:`blocks <dataset_concept>`.
 
-        When `target_num_rows_per_block` is set, it repartitions :class:`Dataset`
-        to honor target number of rows per :ref:`blocks <dataset_concept>`. Note
-        that the system will internally figure out the number of rows per
-        :ref:`blocks <dataset_concept>` for optimal execution, based on the
-        `target_num_rows_per_block`.
-
         This method can be useful to tune the performance of your pipeline. To learn
         more, see :ref:`Advanced: Performance Tips and Tuning <data_performance_tips>`.
 
@@ -1408,9 +1402,16 @@ def repartition(
 
         Args:
             num_blocks: Number of blocks after repartitioning.
-            target_num_rows_per_block: The target number of rows per block to
+            target_num_rows_per_block: [Experimental] The target number of rows per block to
                 repartition. Note that either `num_blocks` or
-                `target_num_rows_per_block` must be set, but not both.
+                `target_num_rows_per_block` must be set, but not both. When
+                `target_num_rows_per_block` is set, it only repartitions
+                :class:`Dataset` :ref:`blocks <dataset_concept>` that are larger than
+                `target_num_rows_per_block`. Note that the system will internally
+                figure out the number of rows per :ref:`blocks <dataset_concept>` for
+                optimal execution, based on the `target_num_rows_per_block`. This is
+                the current behavior because of the implementation and may change in
+                the future.
             shuffle: Whether to perform a distributed shuffle during the
                 repartition. When shuffle is enabled, each output block
                 contains a subset of data rows from each input block, which

python/ray/data/random_access_dataset.py

@@ -234,8 +234,10 @@ def multiget(self, block_indices, keys):
             col = block[self.key_field]
             indices = np.searchsorted(col, keys)
             acc = BlockAccessor.for_block(block)
-            result = [acc._get_row(i) for i in indices]
-            # assert result == [self._get(i, k) for i, k in zip(block_indices, keys)]
+            result = [
+                acc._get_row(i) if k1.as_py() == k2 else None
+                for i, k1, k2 in zip(indices, col.take(indices), keys)
+            ]
         else:
             result = [self._get(i, k) for i, k in zip(block_indices, keys)]
         self.total_time += time.perf_counter() - start