ronantakizawa/codereview-bench huggingface dataset:
Repository: executablebooks/jupyter-book
File: jupyter_book/page.py (Python)

Code before:
    if execute is True:
        ntbk = run_ntbk(ntbk, op.dirname(path_ntbk))

    # Define the path to images and then the relative path to where they'll originally be placed
    if isinstance(path_media_output, str):
        path_media_output_rel = op.relpath(path_media_output, path_html_output)

    # Generate HTML from our notebook using the template
    output_resources = {'output_files_dir': path_media_output_rel, 'unique_key': notebook_name}
    exp = HTMLExporter(template_file=path_template, config=c)
    html, resources = exp.from_notebook_node(ntbk, resources=output_resources)

    # Now write the markdown and resources
    writer = FilesWriter(config=c)
    writer.write(html, resources, notebook_name=notebook_name)
    if verbose:
        print("Finished writing notebook to {}".format(path_html_output))


def _is_jupytext_file(ntbk):
    """Infer whether a notebook node was created from a Jupytext Markdown file.

    Right now, this just tries to guess based on whether there's a particular piece of
    metadata in the notebook. Not sure if this is actually correct though...
    """
    return ntbk.get('metadata', {}).get('jupytext', {}).get('cell_metadata_filter', '') != "-all"


Review comment (suggestion):
I think you want to test that the metadata `notebook_metadata_filter` is `-all`. I will explain this with a bit more detail at https://github.com/mwouts/jupytext/issues/320

Code after:
    # Excution of the notebook if we wish
    if execute is True:
        ntbk = run_ntbk(ntbk, op.dirname(path_ntbk))

    # Define the path to images and then the relative path to where they'll originally be placed
    if isinstance(path_media_output, str):
        path_media_output_rel = op.relpath(path_media_output, path_html_output)

    # Generate HTML from our notebook using the template
    output_resources = {'output_files_dir': path_media_output_rel, 'unique_key': notebook_name}
    exp = HTMLExporter(template_file=path_template, config=c)
    html, resources = exp.from_notebook_node(ntbk, resources=output_resources)

    # Now write the markdown and resources
    writer = FilesWriter(config=c)
    writer.write(html, resources, notebook_name=notebook_name)
    if verbose:
        print("Finished writing notebook to {}".format(path_html_output))


def _is_jupytext_file(ntbk):
    """Infer whether a notebook node was created from a Jupytext Markdown file.

    Right now, this just tries to guess based on whether there's a particular piece of
    metadata in the notebook.
    """
    jupytext_meta = ntbk.get('metadata', {}).get('jupytext')
    if jupytext_meta is None:
        return False
    else:
        return jupytext_meta.get('notebook_metadata_filter', '') != "-all"
Repository: executablebooks/jupyter-book
File: tests/test_tocdirective.py (Python)

Code before:
    path_output = Path(tmpdir).joinpath("mybook").absolute()
    # Regular TOC should work
    p_toc = path_books.joinpath("toc")
    path_toc = p_toc.joinpath("_toc_withheaders.yml")
    run(
        f"jb build {p_toc} --path-output {path_output} --toc {path_toc} -W".split(),
        check=True,
    )

    path_toc_directive = path_output.joinpath("_build", "html", "index.html")
    file_regression.check(path_toc_directive.read_text(), extension=".html")


def test_toc_url(tmpdir, file_regression):
    """Testing with additional `url` key in _toc.yml"""
    path_output = Path(tmpdir).joinpath("mybook").absolute()
    # Regular TOC should work
    p_toc = path_books.joinpath("toc")
    path_toc = p_toc.joinpath("_toc_url.yml")
    run(
        f"jb build {p_toc} --path-output {path_output} --toc {path_toc} -W".split(),
        check=True,
    )

    path_toc_directive = path_output.joinpath("_build", "html", "index.html")
    file_regression.check(path_toc_directive.read_text(), extension=".html")


Review comment (suggestion):
here could you use BeautifulSoup or something like this to *only* pull out the HTML of the toc block that is created by the directive? Otherwise we are gonna error any time the header changes in sphinx-book-theme etc

Code after:
    file_regression.check(str(toc), extension=".html")


def test_toc_withheaders(tmpdir, file_regression):
    """Testing `header` in _toc.yml"""
    path_output = Path(tmpdir).joinpath("mybook").absolute()
    # Regular TOC should work
    p_toc = path_books.joinpath("toc")
    path_toc = p_toc.joinpath("_toc_withheaders.yml")
    run(
        f"jb build {p_toc} --path-output {path_output} --toc {path_toc} -W".split(),
        check=True,
    )

    path_toc_directive = path_output.joinpath("_build", "html", "index.html")

    # get the tableofcontents markup
    soup = BeautifulSoup(path_toc_directive.read_text(), "html.parser")
    toc = soup.find_all("div", class_="tableofcontents-wrapper")[0]

    file_regression.check(str(toc), extension=".html")


def test_toc_urllink(tmpdir, file_regression):
    """Testing with additional `url` link key in _toc.yml"""
    path_output = Path(tmpdir).joinpath("mybook").absolute()
    # Regular TOC should work
    p_toc = path_books.joinpath("toc")
    path_toc = p_toc.joinpath("_toc_urllink.yml")
    run(
        f"jb build {p_toc} --path-output {path_output} --toc {path_toc} -W".split(),
        check=True,
    )

    path_toc_directive = path_output.joinpath("_build", "html", "index.html")

    # get the tableofcontents markup
    soup = BeautifulSoup(path_toc_directive.read_text(), "html.parser")
    toc = soup.find_all("div", class_="tableofcontents-wrapper")[0]
    file_regression.check(str(toc), extension=".html")
Repository: executablebooks/jupyter-book
File: jupyter_book/config.py (Python)

Code before:
            if isinstance(val, list) and list_extend:
                config[key].extend(val)
            else:
                config[key] = val
        else:
            config[key] = val


def _get_files_outside_toc(
    toc: Path, sourcedir: Path, excluded_patterns: Collection[str]
):
    source_root = sourcedir or Path()
    source_files = {f for f in glob(str(source_root / "**/*"), recursive=True)}

    excluded_file_sets = [set(glob(p, recursive=True)) for p in excluded_patterns]
    included_files: Set[str] = {
        Path(relpath(f, source_root)).as_posix()
        for f in source_files.difference(*excluded_file_sets)
        if not isdir(f)
    }

    toc_yaml = yaml.safe_load(toc.read_text(encoding="utf8"))

    from nested_lookup import nested_lookup

    toc_files = {f for f in nested_lookup("file", toc_yaml)}

    verified_toc_files: Set[str] = {
        Path(f).as_posix()
        for f in included_files
        if os.path.splitext(f)[0] in toc_files
    }
    return included_files.difference(verified_toc_files)


Review comment (suggestion):
nested lookup is cool! I guess it is probably fairly lightweight? In general we want to be conservative about adding dependencies - do you think the extra dep is worth the functionality?

(note that if we do add the dep, there are probably other places in the codebase where it'd be useful as we are doing some recursive stuff here and there I think)

Code after:
        elif isinstance(config.get(key), list):
            if isinstance(val, list) and list_extend:
                config[key].extend(val)
            else:
                config[key] = val
        else:
            config[key] = val


def _get_files_outside_toc(
    toc: Path, sourcedir: Path, excluded_patterns: Collection[str]
) -> Set[str]:
    """Returns a set of files that are outside of the toc for exclusion from sphinx.

    Hidden files are NOT processed here as it may result in thousands of individual
     exclusions.
    """
    source_root = sourcedir or Path()
    source_files = {ff for ff in glob(str(source_root / "**/*"), recursive=True)}

    excluded_file_sets = [set(glob(pp, recursive=True)) for pp in excluded_patterns]
    included_files: Set[str] = {
        Path(relpath(ff, source_root)).as_posix()
        for ff in source_files.difference(*excluded_file_sets)
        if not isdir(ff)
    }

    toc_yaml = yaml.safe_load(toc.read_text(encoding="utf8"))
    toc_files = {ff for ff in nested_lookup("file", toc_yaml)}

    verified_toc_files: Set[str] = {
        Path(ff).as_posix()
        for ff in included_files
        if os.path.splitext(ff)[0] in toc_files
    }
    return included_files.difference(verified_toc_files)
Repository: exo-explore/exo
File: src/exo/worker/utils/net_profile.py (Python)

Code before:
import http.client

from anyio import create_task_group, to_thread
from loguru import logger

from exo.shared.topology import Topology
from exo.shared.types.common import NodeId


async def check_reachability(
    target_ip: str,
    expected_node_id: NodeId,
    self_node_id: NodeId,
    out: dict[NodeId, set[str]],
) -> None:
    """Check if a node is reachable at the given IP and verify its identity."""

    def _fetch_remote_node_id() -> str | None:
        connection = http.client.HTTPConnection(target_ip, 52415, timeout=1)
        try:
            connection.request("GET", "/node_id")
            response = connection.getresponse()
            if response.status != 200:
                return None

            body = response.read().decode("utf-8").strip()

            # Strip quotes if present (JSON string response)
            if body.startswith('"') and body.endswith('"') and len(body) >= 2:
                body = body[1:-1]

            return body or None
        except OSError:
            return None
        finally:

Review comment (suggestion):
I know @Evanev7 had a different solution for this.
Here we're hitting the `/node_id` endpoint on the HTTP API on each node.
Evan's solution was lower level, I think checking if it's possible to open a TCP socket.
However, the tradeoff is that Evan's solution ignores loopback addresses which makes it so this ping wouldn't work for a second node that's running on the same machine. This is an annoying constraint if we can't run two exo's on on machine (super helpful for testing). However, maybe discovery would still work fine. It's a question for @Evanev7 

Code after:
import http.client

from anyio import create_task_group, to_thread
from loguru import logger

from exo.shared.topology import Topology
from exo.shared.types.common import NodeId


async def check_reachability(
    target_ip: str,
    expected_node_id: NodeId,
    self_node_id: NodeId,
    out: dict[NodeId, set[str]],
) -> None:
    """Check if a node is reachable at the given IP and verify its identity."""

    def _fetch_remote_node_id() -> NodeId | None:
        connection = http.client.HTTPConnection(target_ip, 52415, timeout=1)
        try:
            connection.request("GET", "/node_id")
            response = connection.getresponse()
            if response.status != 200:
                return None

            body = response.read().decode("utf-8").strip()

            # Strip quotes if present (JSON string response)
            if body.startswith('"') and body.endswith('"') and len(body) >= 2:
                body = body[1:-1]

            return NodeId(body) or None
        except OSError:
            return None
        finally:
Repository: exo-explore/exo
File: app/EXO/EXO/main.swift (Swift)

Code before:
//
//  main.swift
//  EXO
//
//  Created by Jake Hillion on 2025-02-03.
//

import Foundation

/// Command line options for the EXO app
enum CLICommand {
    case install
    case uninstall
    case help
    case none
}

/// Parse command line arguments to determine the CLI command
func parseArguments() -> CLICommand {
    let args = CommandLine.arguments
    if args.contains("--help") || args.contains("-h") {
        return .help
    }
    if args.contains("--install") {
        return .install
    }
    if args.contains("--uninstall") {
        return .uninstall
    }
    return .none

Review comment (suggestion):
```suggestion
//  Created by Jake Hillion on 2026-02-03.
```

Code after:
//
//  main.swift
//  EXO
//
//  Created by Jake Hillion on 2026-02-03.
//

import Foundation

/// Command line options for the EXO app
enum CLICommand {
    case install
    case uninstall
    case help
    case none
}

/// Parse command line arguments to determine the CLI command
func parseArguments() -> CLICommand {
    let args = CommandLine.arguments
    if args.contains("--help") || args.contains("-h") {
        return .help
    }
    if args.contains("--install") {
        return .install
    }
    if args.contains("--uninstall") {
        return .uninstall
    }
    return .none
Repository: exo-explore/exo
File: src/exo/worker/engines/mlx/utils_mlx.py (Python)

Code before:
from exo.shared.types.common import Host
from exo.shared.types.memory import Memory
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import (
    BoundInstance,
    MlxJacclInstance,
    MlxRingInstance,
)
from exo.shared.types.worker.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
    ShardMetadata,
    TensorShardMetadata,
)
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.auto_parallel import (
    TimeoutCallback,
    eval_with_timeout,
    pipeline_auto_parallel,
    tensor_auto_parallel,
)
from exo.worker.runner.bootstrap import logger

Group = mx.distributed.Group
# Needed for 8 bit model
_soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(_soft, 2048), _hard), _hard))


# TODO: Test this
#  ALSO https://github.com/exo-explore/exo/pull/233#discussion_r2549683673
def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
    return Memory.from_float_kb(
        (model_shard_meta.end_layer - model_shard_meta.start_layer)
        / model_shard_meta.n_layers
        * model_shard_meta.model_card.storage_size.in_kb
        / (
            1
            if isinstance(model_shard_meta, PipelineShardMetadata)
            else model_shard_meta.world_size
        )
    )


class ModelLoadingTimeoutError(Exception):
    pass


def mx_barrier(group: Group | None = None):
    mx.eval(
        mx.distributed.all_sum(

Review comment (refactor):
Instead of extending this code in the globals, could you move this into `main` in `runner.py`? Thanks!

Code after:
from exo.shared.types.memory import Memory
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import (
    BoundInstance,
    MlxJacclInstance,
    MlxRingInstance,
)
from exo.shared.types.worker.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
    ShardMetadata,
    TensorShardMetadata,
)
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.auto_parallel import (
    TimeoutCallback,
    eval_with_timeout,
    pipeline_auto_parallel,
    tensor_auto_parallel,
)
from exo.worker.runner.bootstrap import logger

Group = mx.distributed.Group


# TODO: Test this
#  ALSO https://github.com/exo-explore/exo/pull/233#discussion_r2549683673
def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
    return Memory.from_float_kb(
        (model_shard_meta.end_layer - model_shard_meta.start_layer)
        / model_shard_meta.n_layers
        * model_shard_meta.model_card.storage_size.in_kb
        / (
            1
            if isinstance(model_shard_meta, PipelineShardMetadata)
            else model_shard_meta.world_size
        )
    )


class ModelLoadingTimeoutError(Exception):
    pass


def mx_barrier(group: Group | None = None):
    mx.eval(
        mx.distributed.all_sum(
            mx.array(1.0),
            stream=mx.default_stream(mx.Device(mx.cpu)),
            group=group,
        )
Repository: exo-explore/exo
File: rust/exo_pyo3_bindings/src/networking.rs (Rust)

Code before:
        fn __str__(&self) -> String {
            Self::MSG.to_string()
        }
    }
}

/// Connection or disconnection event discriminant type.
#[gen_stub_pyclass_enum]
#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
#[derive(Debug, Clone, PartialEq)]
enum PyConnectionUpdateType {
    Connected = 0,
    Disconnected,
}

#[gen_stub_pyclass]
#[pyclass(frozen, name = "ConnectionUpdate")]
#[derive(Debug, Clone)]
struct PyConnectionUpdate {
    /// Whether this is a connection or disconnection event
    #[pyo3(get)]
    update_type: PyConnectionUpdateType,

    /// Identity of the peer that we have connected to or disconnected from.
    #[pyo3(get)]
    peer_id: PyPeerId,

    /// Remote connection's IPv4 address.
    #[pyo3(get)]
    remote_ipv4: String,

    /// Remote connection's TCP port.
    #[pyo3(get)]
    remote_tcp_port: u16,
}

enum ToTask {
    GossipsubSubscribe {
        topic: String,
        result_tx: oneshot::Sender<PyResult<bool>>,
    },
    GossipsubUnsubscribe {
        topic: String,
        result_tx: oneshot::Sender<bool>,
    },
    GossipsubPublish {
        topic: String,
        data: Vec<u8>,
        result_tx: oneshot::Sender<PyResult<MessageId>>,
    },
}

Review comment (suggestion):
It feels a bit sad to change this interface to a String, was it painful to use before? Might be better to at least have a `struct PeerId(String)`

Code after:
        fn __str__(&self) -> String {
            Self::MSG.to_string()
        }
    }
}

/// Connection or disconnection event discriminant type.
#[gen_stub_pyclass_enum]
#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
#[derive(Debug, Clone, PartialEq)]
enum PyConnectionUpdateType {
    Connected = 0,
    Disconnected,
}

#[gen_stub_pyclass]
#[pyclass(frozen, name = "ConnectionUpdate")]
#[derive(Debug, Clone)]
struct PyConnectionUpdate {
    /// Whether this is a connection or disconnection event
    #[pyo3(get)]
    update_type: PyConnectionUpdateType,

    /// Identity of the peer that we have connected to or disconnected from.
    #[pyo3(get)]
    peer_id: String,

    /// Remote connection's IPv4 address.
    #[pyo3(get)]
    remote_ipv4: String,

    /// Remote connection's TCP port.
    #[pyo3(get)]
    remote_tcp_port: u16,
}

enum ToTask {
    GossipsubSubscribe {
        topic: String,
        result_tx: oneshot::Sender<PyResult<bool>>,
    },
    GossipsubUnsubscribe {
        topic: String,
        result_tx: oneshot::Sender<bool>,
    },
    GossipsubPublish {
        topic: String,
        data: Vec<u8>,
        result_tx: oneshot::Sender<PyResult<MessageId>>,
    },
}
Repository: exo-explore/exo
File: src/exo/utils/info_gatherer/net_profile.py (Python)

Code before:
    """Yield (ip, node_id) pairs as reachability probes complete."""

    send, recv = channel[tuple[str, NodeId]]()

    # these are intentionally httpx's defaults so we can tune them later
    timeout = httpx.Timeout(timeout=5.0)
    limits = httpx.Limits(
        max_connections=100,
        max_keepalive_connections=20,
        keepalive_expiry=5,
    )

    async def _probe(
        target_ip: str,
        expected_node_id: NodeId,
        client: httpx.AsyncClient,
        send: Sender[tuple[str, NodeId]],
    ) -> None:
        async with send:
            out: defaultdict[NodeId, set[str]] = defaultdict(set)
            await check_reachability(target_ip, expected_node_id, out, client)
            if expected_node_id in out:
                await send.send((target_ip, expected_node_id))

    async with (
        httpx.AsyncClient(timeout=timeout, limits=limits) as client,
        create_task_group() as tg,
    ):
        for node_id in topology.list_nodes():
            if node_id not in node_network:
                continue
            if node_id == self_node_id:
                continue
            for iface in node_network[node_id].interfaces:
                tg.start_soon(_probe, iface.ip_address, node_id, client, send.clone())
        send.close()

        with recv:
            async for item in recv:
                yield item


Review comment (suggestion):
This is pretty suspicious, it should be a no-op for `http://` URLs which are what we should be using.

Code after:
    """Yield (ip, node_id) pairs as reachability probes complete."""

    send, recv = channel[tuple[str, NodeId]]()

    # these are intentionally httpx's defaults so we can tune them later
    timeout = httpx.Timeout(timeout=5.0)
    limits = httpx.Limits(
        max_connections=100,
        max_keepalive_connections=20,
        keepalive_expiry=5,
    )

    async def _probe(
        target_ip: str,
        expected_node_id: NodeId,
        client: httpx.AsyncClient,
        send: Sender[tuple[str, NodeId]],
    ) -> None:
        async with send:
            out: defaultdict[NodeId, set[str]] = defaultdict(set)
            await check_reachability(target_ip, expected_node_id, out, client)
            if expected_node_id in out:
                await send.send((target_ip, expected_node_id))

    async with (
        httpx.AsyncClient(timeout=timeout, limits=limits, verify=False) as client,
        create_task_group() as tg,
    ):
        for node_id in topology.list_nodes():
            if node_id not in node_network:
                continue
            if node_id == self_node_id:
                continue
            for iface in node_network[node_id].interfaces:
                tg.start_soon(_probe, iface.ip_address, node_id, client, send.clone())
        send.close()

        with recv:
            async for item in recv:
                yield item