ronantakizawa/codereview-bench huggingface dataset:
Repository: ray-project/ray
File: python/ray/train/v2/_internal/execution/callback_manager.py (Python)

Code before:
import logging

from ray.train.v2._internal.execution.callback import CallbackErrorAction
from ray.train.v2.api.exceptions import ControllerError, TrainingFailedError

logger = logging.getLogger(__name__)


class CallbackManager:
    def __init__(self, callbacks):
        self._callbacks = callbacks

    # change this return type later
    def invoke(self, hook_name: str, *args, **context) -> TrainingFailedError | None:
        for callback in self._callbacks:
            callback_name = type(callback).__name__
            method = getattr(callback, hook_name, None)
            if not method:
                continue
            try:
                method(*args, **context)
            except Exception as e:
                exc_handler = getattr(callback, "on_callback_hook_exception", None)
                if exc_handler is None:
                    continue
                try:
                    result = exc_handler(hook_name, e, **context)
                except Exception as handler_exc:
                    logger.exception(
                        f"Exception raised in callback hook 'on_callback_hook_exception' "
                        f"from callback '{callback_name}' while handling hook '{hook_name}'"
                    )
                    return ControllerError(handler_exc)

                if not self._is_callback_hook_exception_handler_result(result):
                    e = TypeError(
                        "`on_callback_hook_exception` must return "
                        "(CallbackErrorAction, Optional[TrainingFailedError]), "
                        f"got {type(result)}"

Review comment (suggestion):
Maybe we should replace hook_name with an enum so that there is a clear contract on the possible hooks- less error prone than a raw string. Could have a different enum for worker group and controller hooks. 

Code after:
import logging

from ray.train.v2.api.exceptions import ControllerError

logger = logging.getLogger(__name__)


class CallbackManager:
    def __init__(self, callbacks):
        self._callbacks = callbacks

    # change this return type later
    def invoke(self, hook_name: str, *args, **context) -> None:
        for callback in self._callbacks:
            callback_name = type(callback).__name__
            method = getattr(callback, hook_name, None)
            if method is None or not callable(method):
                raise ControllerError(
                    AttributeError(
                        f"Callback '{callback_name}' hook '{hook_name}' is missing "
                        "or not callable."
                    )
                )
            try:
                method(*args, **context)
            except Exception as e:
                # TODO: Enable configuration to suppress exceptions.
                logger.exception(
                    f"Exception raised in callback hook '{hook_name}' from callback "
                    f"'{callback_name}'."
                )
                raise ControllerError(e) from e
Repository: ray-project/ray
File: python/ray/train/v2/_internal/execution/callback_manager.py (Python)

Code before:
                        "`on_callback_hook_exception` must return "
                        "(CallbackErrorAction, Optional[TrainingFailedError]), "
                        f"got {type(result)}"
                    )
                    return ControllerError(e)

                action, mapped_error = result
                if not mapped_error:
                    continue

                match action:
                    case CallbackErrorAction.SUPPRESS:
                        if mapped_error is not None:
                            logger.exception(
                                f"Exception raised in callback hook '{hook_name}' from callback '{callback_name}'."
                            )
                        continue
                    case CallbackErrorAction.RAISE:
                        return mapped_error
                    case _:
                        e = ValueError(f"Unknown CallbackErrorAction: {action}")
                        return ControllerError(e)

        return None

    def _is_callback_hook_exception_handler_result(self, result: object) -> bool:
        if not (isinstance(result, tuple) and len(result) == 2):
            return False
        action, mapped_error = result
        if not isinstance(action, CallbackErrorAction):
            return False
        if mapped_error is not None and not isinstance(
            mapped_error, TrainingFailedError
        ):
            return False
        return True


Review comment (suggestion):
rename to something simpler like _validate_handler_result

Code after:
import logging

from ray.train.v2.api.exceptions import ControllerError

logger = logging.getLogger(__name__)


class CallbackManager:
    def __init__(self, callbacks):
        self._callbacks = callbacks

    # change this return type later
    def invoke(self, hook_name: str, *args, **context) -> None:
        for callback in self._callbacks:
            callback_name = type(callback).__name__
            method = getattr(callback, hook_name, None)
            if method is None or not callable(method):
                raise ControllerError(
                    AttributeError(
                        f"Callback '{callback_name}' hook '{hook_name}' is missing "
                        "or not callable."
                    )
                )
            try:
                method(*args, **context)
            except Exception as e:
                # TODO: Enable configuration to suppress exceptions.
                logger.exception(
                    f"Exception raised in callback hook '{hook_name}' from callback "
                    f"'{callback_name}'."
                )
                raise ControllerError(e) from e
Repository: ray-project/ray
File: src/ray/protobuf/public/events_actor_lifecycle_event.proto (Protocol Buffers)

Code before:
    // Actor is being restarted because the node was preempted.
    NODE_PREEMPTION = 2;
  }

  message StateTransition {
    State state = 1;
    google.protobuf.Timestamp timestamp = 2;
    // The node id of the actor once it is created.
    // available when state is ALIVE updated when the actor is restarted.
    bytes node_id = 3;
    // The worker id of the worker on which this actor is running. available when state is ALIVE.
    // The worker id can change when the actor is restarted.
    bytes worker_id = 4;
    // The repr name of the actor if specified with a customized repr method, e.g. __repr__
    // This field is only available after the actor creation task has been run since it
    // might depend on actor fields to be initialized in __init__.
    // Default to empty string if no customized repr is defined.
    string repr_name = 5;
    // Contains metadata about why the actor is dead. available when state is DEAD.
    ActorDeathCause death_cause = 6;
    // PID of the actor process. Available when state is ALIVE.
    uint32 pid = 7;
    // Port of the actor worker. Available when state is ALIVE.
    int32 port = 8;
    // Reason for restarting. Only set when state is RESTARTING.
    RestartReason restarting_reason = 9;
  }

  // The ID of the actor that was created.
  bytes actor_id = 1;
  // Current state of this actor.
  repeated StateTransition state_transitions = 2;
}


Review comment (suggestion):
align the name with the type name: either `RestartingReason` or `restart_reason`

Code after:
    // Actor is being restarted because the node was preempted.
    NODE_PREEMPTION = 2;
  }

  message StateTransition {
    State state = 1;
    google.protobuf.Timestamp timestamp = 2;
    // The node id of the actor once it is created.
    // available when state is ALIVE updated when the actor is restarted.
    bytes node_id = 3;
    // The worker id of the worker on which this actor is running. available when state is ALIVE.
    // The worker id can change when the actor is restarted.
    bytes worker_id = 4;
    // The repr name of the actor if specified with a customized repr method, e.g. __repr__
    // This field is only available after the actor creation task has been run since it
    // might depend on actor fields to be initialized in __init__.
    // Default to empty string if no customized repr is defined.
    string repr_name = 5;
    // Contains metadata about why the actor is dead. available when state is DEAD.
    ActorDeathCause death_cause = 6;
    // PID of the actor process. Available when state is ALIVE.
    uint32 pid = 7;
    // Port of the actor worker. Available when state is ALIVE.
    int32 port = 8;
    // Reason for restarting. Only set when state is RESTARTING.
    RestartReason restart_reason = 9;
  }

  // The ID of the actor that was created.
  bytes actor_id = 1;
  // Current state of this actor.
  repeated StateTransition state_transitions = 2;
}
Repository: ray-project/ray
File: src/ray/protobuf/public/events_actor_definition_event.proto (Protocol Buffers)

Code before:

message ActorDefinitionEvent {
  // The ID of the actor that was created.
  bytes actor_id = 1;
  // The ID of the job that created the actor.
  bytes job_id = 2;
  // Whether the actor is persistent.
  bool is_detached = 3;
  // Name of the actor.
  string name = 4;
  // The actor's namespace. Named `ray_namespace` to avoid conflicting with c++ keyword.
  string ray_namespace = 5;
  // Serialized runtime_env used to report in the dashboard snapshot. We need to populate
  // it here instead of grabbing it from the task spec because the task spec is cleared
  // for deleted actors: https://github.com/ray-project/ray/pull/11149.
  string serialized_runtime_env = 6;
  // The actor's class name. This is necessary because the task spec's lifetime
  // is shorter than the ActorTableData.
  string class_name = 7;
  // Quantities of the different resources required by this actor.
  map<string, double> required_resources = 8;
  // Placement group ID if the actor requires a placement group.
  bytes placement_group_id = 9;
  // The label selector for the actor.
  map<string, string> label_selector = 11;
  // Human readable stacktrace of the actor creation (from task_spec.call_site).
  optional string call_site = 12;
  // The ID of the caller of the actor creation task (parent actor/driver).
  bytes parent_id = 13;
}


Review comment (security):
cc: @edoakes My understanding is exposing the stacktrace will potentially expose content from user's application. Would we have potential security concerns related to this?

Code after:

message ActorDefinitionEvent {
  // The ID of the actor that was created.
  bytes actor_id = 1;
  // The ID of the job that created the actor.
  bytes job_id = 2;
  // Whether the actor is persistent.
  bool is_detached = 3;
  // Name of the actor.
  string name = 4;
  // The actor's namespace. Named `ray_namespace` to avoid conflicting with c++ keyword.
  string ray_namespace = 5;
  // Serialized runtime_env used to report in the dashboard snapshot. We need to populate
  // it here instead of grabbing it from the task spec because the task spec is cleared
  // for deleted actors: https://github.com/ray-project/ray/pull/11149.
  string serialized_runtime_env = 6;
  // The actor's class name. This is necessary because the task spec's lifetime
  // is shorter than the ActorTableData.
  string class_name = 7;
  // Quantities of the different resources required by this actor.
  map<string, double> required_resources = 8;
  // Placement group ID if the actor requires a placement group.
  bytes placement_group_id = 9;
  // The label selector for the actor.
  map<string, string> label_selector = 11;
  // Human readable stacktrace of the actor creation (from task_spec.call_site).
  string call_site = 12;
  // The ID of the caller of the actor creation task (parent actor/driver).
  bytes parent_id = 13;
}
Repository: ray-project/ray
File: python/ray/data/_internal/logical/rules/__init__.py (Python)

Code before:
"""Expose rule classes in ray.data._internal.logical.rules.

This module dynamically imports each submodule in the package and re-exports
names listed in the submodule's __all__ attribute. This allows callers to
import rules directly from ray.data._internal.logical.rules.
"""

from importlib import import_module
from pkgutil import iter_modules

__all__ = []

for _loader, _module_name, _is_pkg in iter_modules(__path__):
    if _module_name.startswith("_"):
        continue
    _module = import_module(f"{__name__}.{_module_name}")
    _exported_names = getattr(_module, "__all__", [])
    for _name in _exported_names:
        globals()[_name] = getattr(_module, _name)
    __all__.extend(_exported_names)

__all__.sort()


Review comment (performance):
Rather than using a script, can we explicitly list out the dependencies? e.g., `from .inherit_batch_format import ...`? I think it'll be easier to read and maintain

Code after:
"""Expose rule classes in ray.data._internal.logical.rules."""

from .combine_shuffles import CombineShuffles
from .configure_map_task_memory import (
    ConfigureMapTaskMemoryRule,
    ConfigureMapTaskMemoryUsingOutputSize,
)
from .inherit_batch_format import InheritBatchFormatRule
from .inherit_target_max_block_size import InheritTargetMaxBlockSizeRule
from .limit_pushdown import LimitPushdownRule
from .operator_fusion import FuseOperators, are_remote_args_compatible
from .predicate_pushdown import PredicatePushdown
from .projection_pushdown import ProjectionPushdown
from .set_read_parallelism import (
    SetReadParallelismRule,
    compute_additional_split_factor,
)

__all__ = [
    "CombineShuffles",
    "ConfigureMapTaskMemoryRule",
    "ConfigureMapTaskMemoryUsingOutputSize",
    "FuseOperators",
    "InheritBatchFormatRule",
    "InheritTargetMaxBlockSizeRule",
    "LimitPushdownRule",
    "PredicatePushdown",
    "ProjectionPushdown",
    "SetReadParallelismRule",
    "are_remote_args_compatible",
    "compute_additional_split_factor",
]
Repository: ray-project/ray
File: python/ray/llm/tests/batch/cpu/stages/test_http_request_stage.py (Python)

Code before:
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": "http://example.com/image.png"},
                            },
                            {
                                "type": "text",
                                "text": "hello",
                            },
                        ],
                    }
                ],
            }
        }
    ]
    config = HttpRequestProcessorConfig(
        url=str(numpy_payload_server.make_url("/")),
        headers={"Content-Type": "application/json"},
        qps=None,
    )

    # construct udf and invoke http_request_stage in map_batches.
    # Note that default batch_format in map_batches is numpy.
    processor = ProcessorBuilder.build(config)
    ds = processor(ray.data.from_items(data * 10))
    results = await asyncio.to_thread(ds.take_all)

    # Verify results
    assert len(results) == 10
    for result in results:
        assert result["http_response"]["response"] == "success"


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", __file__]))


Review comment (suggestion):
I don't think we need to place this on an asyncio thread. We can invoke `results = ds.take_all()` directly.

Code after:
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": "http://example.com/image.png"},
                            },
                            {
                                "type": "text",
                                "text": "hello",
                            },
                        ],
                    }
                ],
            }
        }
    ]
    config = HttpRequestProcessorConfig(
        url=str(numpy_payload_server.make_url("/")),
        headers={"Content-Type": "application/json"},
        qps=None,
    )

    processor = ProcessorBuilder.build(config)
    ds = processor(ray.data.from_items(data * 10))
    results = await asyncio.to_thread(ds.take_all)

    assert len(results) == 10
    for result in results:
        assert result["http_response"]["response"] == "success"


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", __file__]))
Repository: ray-project/ray
File: python/ray/data/_internal/util.py (Python)

Code before:
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generator,
    Iterable,
    Iterator,
    List,
    Optional,
    Tuple,
    TypeVar,
    Union,
)

import numpy as np
import pandas as pd

# NOTE: pyarrow.fs module needs to be explicitly imported!
import pyarrow
import pyarrow.fs

import ray
from ray._common.retry import call_with_retry
from ray.data.context import DEFAULT_READ_OP_MIN_NUM_BLOCKS, WARN_PREFIX, DataContext
from ray.data.util.arrow_utils import (
    _check_pyarrow_version as arrow_utils_check_pyarrow_version,
)
from ray.util.annotations import DeveloperAPI

import psutil

if TYPE_CHECKING:
    import pandas

    from ray.data._internal.compute import ComputeStrategy
    from ray.data._internal.execution.interfaces import RefBundle
    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
    from ray.data.block import (
        Block,
        BlockMetadataWithSchema,
        Schema,
        UserDefinedFunction,
    )
    from ray.data.datasource import Datasource, Reader
    from ray.util.placement_group import PlacementGroup

logger = logging.getLogger(__name__)


KiB = 1024  # bytes
MiB = 1024 * KiB

Review comment (suggestion):
If `_check_pyarrow_version` is only used internally, it should still have the underscore prefix to follow "pythonic" practices.

Code after:
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generator,
    Iterable,
    Iterator,
    List,
    Optional,
    Tuple,
    TypeVar,
    Union,
)

import numpy as np
import pandas as pd

# NOTE: pyarrow.fs module needs to be explicitly imported!
import pyarrow
import pyarrow.fs

import ray
from ray._common.retry import call_with_retry
from ray.data.context import DEFAULT_READ_OP_MIN_NUM_BLOCKS, WARN_PREFIX, DataContext
from ray.util.annotations import DeveloperAPI

import psutil

if TYPE_CHECKING:
    import pandas

    from ray.data._internal.compute import ComputeStrategy
    from ray.data._internal.execution.interfaces import RefBundle
    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
    from ray.data.block import (
        Block,
        BlockMetadataWithSchema,
        Schema,
        UserDefinedFunction,
    )
    from ray.data.datasource import Datasource, Reader
    from ray.util.placement_group import PlacementGroup

logger = logging.getLogger(__name__)


KiB = 1024  # bytes
MiB = 1024 * KiB
GiB = 1024 * MiB
Repository: ray-project/ray
File: doc/source/data/doc_code/working-with-llms/basic_llm_example.py (Python)

Code before:
    os.makedirs(path)

# __row_level_fault_tolerance_config_example_start__
# Row-level fault tolerance configuration
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    concurrency=1,
    batch_size=64,
    should_continue_on_error=True,
)
# __row_level_fault_tolerance_config_example_end__

# __checkpoint_config_setup_example_start__
from ray.data.checkpoint import CheckpointConfig

ctx = ray.data.DataContext.get_current()
ctx.checkpoint_config = CheckpointConfig(
    id_column="id",
    checkpoint_path=checkpoint_path,
    delete_checkpoint_on_success=False,
)
# __checkpoint_config_setup_example_end__

# __checkpoint_usage_example_start__
processor_config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen3-0.6B",
    concurrency=1,
    batch_size=16,
    chat_template_stage=False,
    tokenize_stage=True,
    detokenize_stage=False,
)

processor = build_processor(
    processor_config,
    preprocess=lambda row: dict(
        id=row["id"], # Preserve the ID column for checkpointing
        prompt=row["prompt"],
        sampling_params=dict(
            temperature=0.3,
            max_tokens=10,
        ),
    ),
    postprocess=lambda row: {
        "id": row["id"], # Preserve the ID column for checkpointing
        "answer": row.get("generated_text"),
    },
)

ds = ray.data.read_parquet(input_path)
ds = processor(ds)

Review comment (suggestion):
```suggestion
    model_source="unsloth/Llama-3.1-8B-Instruct",
```

Code after:
    os.makedirs(path)

# __row_level_fault_tolerance_config_example_start__
# Row-level fault tolerance configuration
config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    concurrency=1,
    batch_size=64,
    should_continue_on_error=True,
)
# __row_level_fault_tolerance_config_example_end__

# __checkpoint_config_setup_example_start__
from ray.data.checkpoint import CheckpointConfig

ctx = ray.data.DataContext.get_current()
ctx.checkpoint_config = CheckpointConfig(
    id_column="id",
    checkpoint_path=checkpoint_path,
    delete_checkpoint_on_success=False,
)
# __checkpoint_config_setup_example_end__

# __checkpoint_usage_example_start__
processor_config = vLLMEngineProcessorConfig(
    model_source="unsloth/Llama-3.1-8B-Instruct",
    concurrency=1,
    batch_size=16,
)

processor = build_processor(
    processor_config,
    preprocess=lambda row: dict(
        id=row["id"], # Preserve the ID column for checkpointing
        prompt=row["prompt"],
        sampling_params=dict(
            temperature=0.3,
            max_tokens=10,
        ),
    ),
    postprocess=lambda row: {
        "id": row["id"], # Preserve the ID column for checkpointing
        "answer": row.get("generated_text"),
    },
)

ds = ray.data.read_parquet(input_path)
ds = processor(ds)
ds.write_parquet(output_path)
# __checkpoint_usage_example_end__