From a582e5c7031ec9b90ebb06606df15f6fd40851b9 Mon Sep 17 00:00:00 2001
From: Pedro Rodriguez <par@meta.com>
Date: Mon, 12 May 2025 10:00:19 -0700
Subject: [PATCH 1/2] Upgrade stool (#107)

Summary:

Modify stool so that:
- It uses jinja2 for better templating
- Add args to allow force non-interactive overrides
- Add config_parser argument to allow instantiation of pydantic class when some defaults don't exist and use it in stool, to allow the rest of the configuration parser to work.

Test Plan:

```
uv run python -m bytelatent.stool config=/storage/home/par/code/internal-blt/configs/par/stool/run_blt_1b_aws.yaml
uv run python -m bytelatent.train config=../internal-blt/configs/par/train/aws/blt_1b_dclm_aws.yaml dump_dir=/checkpoint/comem/par/scratch/debug data.batch_size=2
```
---
 bytelatent/config_parser.py                  |  19 +++-
 bytelatent/stool.py                          | 113 +++++++------------
 bytelatent/templates/stool_template.sh.jinja |  33 ++++++
 pyproject.toml                               |   1 +
 uv.lock                                      |   4 +-
 5 files changed, 96 insertions(+), 74 deletions(-)
 create mode 100644 bytelatent/templates/stool_template.sh.jinja

diff --git a/bytelatent/config_parser.py b/bytelatent/config_parser.py
index 4c2e2b3..b4bda35 100644
--- a/bytelatent/config_parser.py
+++ b/bytelatent/config_parser.py
@@ -1,9 +1,10 @@
 import copy
-from typing import Type, TypeVar
+from typing import Type, TypeVar, Any
 
 import omegaconf
 from omegaconf import DictConfig, OmegaConf
 from pydantic import BaseModel
+from pydantic_core import PydanticUndefined
 
 
 def parse_file_config(path: str) -> DictConfig:
@@ -68,11 +69,23 @@ def parse_args_with_default(
 
 T = TypeVar("T", bound=BaseModel)
 
+def get_pydantic_default_args(args_cls: Type[T]) -> dict[str, Any]:
+    defaults = {}
+    for field, info in args_cls.model_fields.items():
+        if info.default != PydanticUndefined:
+            defaults[field] = info.default
+    return defaults
 
 def parse_args_to_pydantic_model(
-    args_cls: Type[T], cli_args: DictConfig | None = None
+    args_cls: Type[T], cli_args: DictConfig | None = None, instantiate_default_cls: bool = True
 ) -> T:
-    default_cfg = OmegaConf.create(args_cls().model_dump())
+    if instantiate_default_cls:
+        default_cfg = OmegaConf.create(args_cls().model_dump())
+    else:
+        default_cfg = OmegaConf.create(get_pydantic_default_args(args_cls))
     parsed_cfg = parse_args_with_default(default_cfg=default_cfg, cli_args=cli_args)
+    print(default_cfg)
+    print()
+    print(parsed_cfg)
     pydantic_args = args_cls.model_validate(parsed_cfg)
     return pydantic_args
diff --git a/bytelatent/stool.py b/bytelatent/stool.py
index 90bd2c2..45fb701 100644
--- a/bytelatent/stool.py
+++ b/bytelatent/stool.py
@@ -6,21 +6,28 @@ import shutil
 import subprocess
 from typing import Any, Dict
 
+import jinja2
 from omegaconf import OmegaConf
 from pydantic import BaseModel
 
+from bytelatent.config_parser import parse_args_to_pydantic_model
+
 
 class StoolArgs(BaseModel):
-    name: str = None
-    dump_dir: str = None
-    config: Any = None
+    name: str
+    dump_dir: str
+    # model_config is a reserved name by pydantic, so use this instead
+    model_conf: Any = None
     launcher: str = "sbatch"  # Can be sbatch or bash if already in salloc
+    python_command: str = "python"
+    use_conda: bool = True
     script: str = "apps.main.train"  # The script to run.
     copy_code: bool = True  # Wether to copy code to dump dir
     dirs_exists_ok: bool = (
         False  # Wether to copy new code and config and run regardless that dir exists
     )
-    override: bool = False  # Wether to delete dump dir and restart
+    override: bool = False  # Whether to delete dump dir and restart, requires confirmation
+    force_override: bool = False # Does not require interaction
     nodes: int = -1  # The number of nodes to run the job on.
     ngpu: int = 8  # The number of GPUs required per node.
     ncpu: int = 16  # The number of CPUs allocated per GPU.
@@ -33,41 +40,9 @@ class StoolArgs(BaseModel):
     qos: str = ""
     partition: str = "learn"
     stdout: bool = False
+    dry_run: bool = False
 
 
-SBATCH_COMMAND = """#!/bin/bash
-
-{exclude}
-{qos}
-{account}
-{constraint}
-#SBATCH --job-name={name}
-#SBATCH --nodes={nodes}
-#SBATCH --gres=gpu:{ngpus}
-#SBATCH --cpus-per-gpu={ncpu}
-#SBATCH --time={time}
-#SBATCH --partition={partition}
-#SBATCH --mem={mem}
-
-#SBATCH --output={dump_dir}/logs/%j/%j.stdout
-#SBATCH --error={dump_dir}/logs/%j/%j.stderr
-
-#SBATCH --open-mode=append
-#SBATCH --signal=USR2@120
-#SBATCH --distribution=block
-
-# Mimic the effect of "conda init", which doesn't work for scripts
-eval "$({conda_exe} shell.bash hook)"
-source activate {conda_env_path}
-
-{go_to_code_dir}
-
-export OMP_NUM_THREADS=1
-export LAUNCH_WITH="SBATCH"
-export DUMP_DIR={dump_dir}
-srun {log_output} -n {tasks} -N {nodes_per_run} python -u -m {script} config=$DUMP_DIR/base_config.yaml dump_dir=$DUMP_DIR name={name}
-"""
-
 
 def copy_dir(input_dir: str, output_dir: str) -> None:
     print(f"Copying : {input_dir}\n" f"to      : {output_dir} ...")
@@ -75,6 +50,7 @@ def copy_dir(input_dir: str, output_dir: str) -> None:
     assert os.path.isdir(output_dir), f"{output_dir} is not a directory"
     rsync_cmd = (
         f"rsync -rmt --copy-links "
+        f"--exclude .venv "
         f"--include '**/' "
         f"--include '*.py' "
         f"--exclude='*' "
@@ -151,28 +127,31 @@ def validate_args(args) -> None:
 def launch_job(args: StoolArgs):
     # Set up args default and validate them depending on the cluster or partition requested
     validate_args(args)
-    job_name = args.name or args.config["name"]
-    dump_dir = os.path.join(args.dump_dir, job_name) or args.config["dump_dir"]
+    job_name = args.name or args.model_conf["name"]
+    dump_dir = os.path.join(args.dump_dir, job_name) or args.model_conf["dump_dir"]
     print("Creating directories...")
-    os.makedirs(dump_dir, exist_ok=args.dirs_exists_ok or args.override)
-    if args.override:
-        confirm = input(
-            f"Are you sure you want to delete the directory '{dump_dir}'? This action cannot be undone. (yes/no): "
-        )
-        if confirm.lower() == "yes":
+    os.makedirs(dump_dir, exist_ok=args.dirs_exists_ok or args.override or args.force_override)
+    if args.override or args.force_override:
+        if args.force_override:
             shutil.rmtree(dump_dir)
             print(f"Directory '{dump_dir}' has been deleted.")
         else:
-            print("Operation cancelled.")
-            return
+            confirm = input(
+                f"Are you sure you want to delete the directory '{dump_dir}'? This action cannot be undone. (yes/no): "
+            )
+            if confirm.lower() == "yes":
+                shutil.rmtree(dump_dir)
+                print(f"Directory '{dump_dir}' has been deleted.")
+            else:
+                print("Operation cancelled.")
+                return
     if args.copy_code:
         os.makedirs(f"{dump_dir}/code", exist_ok=args.dirs_exists_ok)
         print("Copying code ...")
         copy_dir(os.getcwd(), f"{dump_dir}/code")
 
     print("Saving config file ...")
-    with open(f"{dump_dir}/base_config.yaml", "w") as cfg:
-        cfg.write(OmegaConf.to_yaml(args.config))
+    shutil.copy(args.model_conf, f"{dump_dir}/base_config.yaml")
 
     conda_exe = os.environ.get("CONDA_EXE", "conda")
     conda_env_path = os.path.dirname(os.path.dirname(args.anaconda))
@@ -181,7 +160,12 @@ def launch_job(args: StoolArgs):
         if not args.stdout
         else ""
     )
-    sbatch = SBATCH_COMMAND.format(
+    env = jinja2.Environment(
+        loader=jinja2.PackageLoader('bytelatent'),
+        autoescape=jinja2.select_autoescape(),
+    )
+    template = env.get_template('stool_template.sh.jinja')
+    sbatch_jinja = template.render(
         name=job_name,
         script=args.script,
         dump_dir=dump_dir,
@@ -197,18 +181,23 @@ def launch_job(args: StoolArgs):
         exclude=args.exclude,
         time=args.time,
         partition=args.partition,
+        python_command=args.python_command,
         conda_exe=conda_exe,
         conda_env_path=conda_env_path,
+        use_conda=args.use_conda,
         log_output=log_output,
         go_to_code_dir=f"cd {dump_dir}/code/" if args.copy_code else "",
     )
 
     print("Writing sbatch command ...")
     with open(f"{dump_dir}/submit.slurm", "w") as f:
-        f.write(sbatch)
+        f.write(sbatch_jinja)
 
-    print("Submitting job ...")
-    os.system(f"{args.launcher} {dump_dir}/submit.slurm")
+    if args.dry_run:
+        print("Dry run mode enabled. Not submitting job.")
+    else:
+        print("Submitting job ...")
+        os.system(f"{args.launcher} {dump_dir}/submit.slurm")
 
     print("Done.")
 
@@ -216,22 +205,6 @@ def launch_job(args: StoolArgs):
 if __name__ == "__main__":
     """
     The command line interface here uses OmegaConf https://omegaconf.readthedocs.io/en/2.3_branch/usage.html#from-command-line-arguments
-    This accepts arguments as a dot list
-    So if the dataclass looks like
-
-    @dataclass
-    class DummyArgs:
-        name: str
-        mode: LMTransformerArgs
-
-    @dataclass
-    class LMTransformerArgs:
-        dim: int
-
-    Then you can pass model.dim=32 to change values in LMTransformerArgs
-    or just name=tictac for top level attributes.
     """
-    args = OmegaConf.from_cli()
-    args.config = OmegaConf.load(args.config)
-    args = StoolArgs.model_validate(args)
+    args = parse_args_to_pydantic_model(StoolArgs, instantiate_default_cls=False)
     launch_job(args)
diff --git a/bytelatent/templates/stool_template.sh.jinja b/bytelatent/templates/stool_template.sh.jinja
new file mode 100644
index 0000000..e6cbd9c
--- /dev/null
+++ b/bytelatent/templates/stool_template.sh.jinja
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+{{ exclude }}
+{{ qos }}
+{{ account }}
+{{ constraint }}
+#SBATCH --job-name={{ name }}
+#SBATCH --nodes={{ nodes }}
+#SBATCH --gres=gpu:{{ ngpus }}
+#SBATCH --cpus-per-gpu={{ ncpu }}
+#SBATCH --time={{ time }}
+#SBATCH --partition={{ partition }}
+#SBATCH --mem={{ mem }}
+
+#SBATCH --output={{ dump_dir }}/logs/%j/%j.stdout
+#SBATCH --error={{ dump_dir }}/logs/%j/%j.stderr
+
+#SBATCH --open-mode=append
+#SBATCH --signal=USR2@120
+#SBATCH --distribution=block
+
+{% if use_conda %}
+# Mimic the effect of "conda init", which doesn't work for scripts
+eval "$({{ conda_exe }} shell.bash hook)"
+source activate {{ conda_env_path }}
+{% endif %}
+
+{{ go_to_code_dir }}
+
+export OMP_NUM_THREADS=1
+export LAUNCH_WITH="SBATCH"
+export DUMP_DIR={{ dump_dir }}
+srun {{ log_output }} -n {{ tasks }} -N {{ nodes_per_run }} {{ python_command }} -u -m {{ script }} config=$DUMP_DIR/base_config.yaml dump_dir=$DUMP_DIR name={{ name }}
diff --git a/pyproject.toml b/pyproject.toml
index 67d86da..ce5cc4e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,6 +9,7 @@ dependencies = [
     "datatrove>=0.5.0",
     "fsspec>=2024.6.1",
     "huggingface-hub==0.30.*",
+    "jinja2>=3.1.6",
     "lm-eval>=0.4.8",
     "luigi>=3.6.0",
     "numpy>=2.1.2",
diff --git a/uv.lock b/uv.lock
index f23aea6..064cb60 100644
--- a/uv.lock
+++ b/uv.lock
@@ -166,6 +166,7 @@ dependencies = [
     { name = "datatrove" },
     { name = "fsspec" },
     { name = "huggingface-hub" },
+    { name = "jinja2" },
     { name = "lm-eval" },
     { name = "luigi" },
     { name = "numpy" },
@@ -201,6 +202,7 @@ requires-dist = [
     { name = "datatrove", specifier = ">=0.5.0" },
     { name = "fsspec", specifier = ">=2024.6.1" },
     { name = "huggingface-hub", specifier = "==0.30.*" },
+    { name = "jinja2", specifier = ">=3.1.6" },
     { name = "lm-eval", specifier = ">=0.4.8" },
     { name = "luigi", specifier = ">=3.6.0" },
     { name = "numpy", specifier = ">=2.1.2" },
@@ -2029,7 +2031,7 @@ wheels = [
 
 [[package]]
 name = "xformers"
-version = "0.0.29+de742ec.d20250502"
+version = "0.0.29+de742ec.d20250507"
 source = { git = "https://github.com/facebookresearch/xformers.git?rev=de742ec3d64bd83b1184cc043e541f15d270c148#de742ec3d64bd83b1184cc043e541f15d270c148" }
 dependencies = [
     { name = "numpy" },

From f95ebd998e3d30a95d4d833a5c23b5cf02f84aec Mon Sep 17 00:00:00 2001
From: Pedro Rodriguez <me@pedro.ai>
Date: Thu, 15 May 2025 22:25:24 +0000
Subject: [PATCH 2/2] Tune retry params as temporary fix for excessive file
 listing

Summary:

Test Plan:
---
 bytelatent/data/file_util.py                |  16 +-
 bytelatent/data/iterators/arrow_iterator.py |  23 ++-
 pyproject.toml                              |   4 +
 uv.lock                                     | 212 ++++++++++++++++++++
 4 files changed, 248 insertions(+), 7 deletions(-)

diff --git a/bytelatent/data/file_util.py b/bytelatent/data/file_util.py
index 5165eab..198cd1a 100644
--- a/bytelatent/data/file_util.py
+++ b/bytelatent/data/file_util.py
@@ -23,12 +23,22 @@ app = typer.Typer()
 S3_PREFIX = "s3://"
 
 
-def get_fs(path: str, s3_profile: str | None = None) -> fsspec.AbstractFileSystem:
+def get_fs(
+    path: str, s3_profile: str | None = None, use_listings_cache: None | bool = None
+) -> fsspec.AbstractFileSystem:
     if path.startswith("s3://"):
+        config_kwargs = {"retries": {"max_attempts": 10, "mode": "standard"}}
         if s3_profile is None:
-            return fsspec.filesystem("s3")
+            return fsspec.filesystem(
+                "s3", use_listings_cache=use_listings_cache, config_kwargs=config_kwargs
+            )
         else:
-            return fsspec.filesystem("s3", profile=s3_profile)
+            return fsspec.filesystem(
+                "s3",
+                profile=s3_profile,
+                use_listings_cache=use_listings_cache,
+                config_kwargs=config_kwargs,
+            )
     else:
         return fsspec.filesystem("file")
 
diff --git a/bytelatent/data/iterators/arrow_iterator.py b/bytelatent/data/iterators/arrow_iterator.py
index ec7c54f..32cdd7f 100644
--- a/bytelatent/data/iterators/arrow_iterator.py
+++ b/bytelatent/data/iterators/arrow_iterator.py
@@ -107,7 +107,13 @@ class ArrowFileIterator(StatefulIterator):
             if self.filesystem_type == "file":
                 self.fs = fsspec.filesystem("file")
             elif self.filesystem_type == "s3":
-                self.fs = fsspec.filesystem("s3", profile=s3_profile)
+                config_kwargs = {"retries": {"max_attempts": 10, "mode": "standard"}}
+                self.fs = fsspec.filesystem(
+                    "s3",
+                    profile=s3_profile,
+                    use_listings_cache=True,
+                    config_kwargs=config_kwargs,
+                )
             else:
                 raise ValueError("Unknown filesystem")
             logger.info("Arrow iterator using fs=%s", self.fs)
@@ -118,7 +124,9 @@ class ArrowFileIterator(StatefulIterator):
             ), "Must specify file_Path if dataset_files is None"
             if file_format == "json":
                 if self.fs is None:
-                    self.fs = get_fs(file_path, s3_profile=s3_profile)
+                    self.fs = get_fs(
+                        file_path, s3_profile=s3_profile, use_listings_cache=True
+                    )
                     if isinstance(self.fs, s3fs.S3FileSystem):
                         self.filesystem_type = "s3"
                     else:
@@ -137,11 +145,16 @@ class ArrowFileIterator(StatefulIterator):
                     data_dir, f"{os.path.basename(jsonl_file)}.shard_*.arrow"
                 )
                 if self.fs is None:
-                    self.fs = get_fs(data_dir_with_glob, s3_profile=s3_profile)
+                    self.fs = get_fs(
+                        data_dir_with_glob,
+                        s3_profile=s3_profile,
+                        use_listings_cache=True,
+                    )
                     if isinstance(self.fs, s3fs.S3FileSystem):
                         self.filesystem_type = "s3"
                     else:
                         self.filesystem_type = "file"
+                logger.info("Globbing: %s", data_dir_with_glob)
                 shard_files = self.fs.glob(data_dir_with_glob)
 
                 for s in shard_files:
@@ -165,7 +178,9 @@ class ArrowFileIterator(StatefulIterator):
                 for f in dataset_files:
                     assert f.startswith("s3://")
             if self.fs is None:
-                self.fs = get_fs(dataset_files[0], s3_profile=s3_profile)
+                self.fs = get_fs(
+                    dataset_files[0], s3_profile=s3_profile, use_listings_cache=True
+                )
                 if isinstance(self.fs, s3fs.S3FileSystem):
                     self.filesystem_type = "s3"
                 else:
diff --git a/pyproject.toml b/pyproject.toml
index ce5cc4e..5fadd7f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,6 +45,10 @@ pre_build = [
     "torch==2.6.0.dev20241112",
 ]
 compile_xformers = ['xformers']
+dev = [
+    "ipython>=9.2.0",
+    "pudb>=2025.1",
+]
 
 
 [tool.uv]
diff --git a/uv.lock b/uv.lock
index 064cb60..e4ab7a7 100644
--- a/uv.lock
+++ b/uv.lock
@@ -148,6 +148,15 @@ version = "4.9.3"
 source = { registry = "https://download.pytorch.org/whl/nightly/cu121" }
 sdist = { url = "https://download.pytorch.org/whl/nightly/antlr4_python3_runtime-4.9.3.tar.gz" }
 
+[[package]]
+name = "asttokens"
+version = "3.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/4a/e7/82da0a03e7ba5141f05cce0d302e6eed121ae055e0456ca228bf693984bc/asttokens-3.0.0.tar.gz", hash = "sha256:0dcd8baa8d62b0c1d118b399b2ddba3c4aff271d0d7a9e0d4c1681c79035bbc7", size = 61978, upload-time = "2024-11-30T04:30:14.439Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/25/8a/c46dcc25341b5bce5472c718902eb3d38600a903b14fa6aeecef3f21a46f/asttokens-3.0.0-py3-none-any.whl", hash = "sha256:e3078351a059199dd5138cb1c706e6430c05eff2ff136af5eb4790f9d28932e2", size = 26918, upload-time = "2024-11-30T04:30:10.946Z" },
+]
+
 [[package]]
 name = "attrs"
 version = "25.3.0"
@@ -190,6 +199,10 @@ dependencies = [
 compile-xformers = [
     { name = "xformers" },
 ]
+dev = [
+    { name = "ipython" },
+    { name = "pudb" },
+]
 pre-build = [
     { name = "ninja" },
     { name = "setuptools" },
@@ -224,6 +237,10 @@ requires-dist = [
 
 [package.metadata.requires-dev]
 compile-xformers = [{ name = "xformers", git = "https://github.com/facebookresearch/xformers.git?rev=de742ec3d64bd83b1184cc043e541f15d270c148" }]
+dev = [
+    { name = "ipython", specifier = ">=9.2.0" },
+    { name = "pudb", specifier = ">=2025.1" },
+]
 pre-build = [
     { name = "ninja" },
     { name = "setuptools" },
@@ -392,6 +409,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/aa/b0/1bdf21c6ff4007aa1ca9566eddd8721b01f398b53a989cdb65cf3026b8f1/datatrove-0.5.0-py3-none-any.whl", hash = "sha256:d65ccafcad4dcc00684b0351734777e01d9a9fef448f57947252591768fef3db", size = 17260206, upload-time = "2025-04-30T16:49:58.62Z" },
 ]
 
+[[package]]
+name = "decorator"
+version = "5.2.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" },
+]
+
 [[package]]
 name = "dill"
 version = "0.3.8"
@@ -434,6 +460,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a2/e7/cbca9e2d2590eb9b5aa8f7ebabe1beb1498f9462d2ecede5c9fd9735faaf/evaluate-0.4.3-py3-none-any.whl", hash = "sha256:47d8770bdea76e2c2ed0d40189273027d1a41ccea861bcc7ba12d30ec5d1e517", size = 84010, upload-time = "2024-09-11T10:15:30.018Z" },
 ]
 
+[[package]]
+name = "executing"
+version = "2.2.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/91/50/a9d80c47ff289c611ff12e63f7c5d13942c65d68125160cefd768c73e6e4/executing-2.2.0.tar.gz", hash = "sha256:5d108c028108fe2551d1a7b2e8b713341e2cb4fc0aa7dcf966fa4327a5226755", size = 978693, upload-time = "2025-01-22T15:41:29.403Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7b/8f/c4d9bafc34ad7ad5d8dc16dd1347ee0e507a52c3adb6bfa8887e1c6a26ba/executing-2.2.0-py2.py3-none-any.whl", hash = "sha256:11387150cad388d62750327a53d3339fad4888b39a6fe233c3afbb54ecffd3aa", size = 26702, upload-time = "2025-01-22T15:41:25.929Z" },
+]
+
 [[package]]
 name = "filelock"
 version = "3.18.0"
@@ -542,6 +577,51 @@ wheels = [
     { url = "https://download.pytorch.org/whl/nightly/idna-3.10-py3-none-any.whl" },
 ]
 
+[[package]]
+name = "ipython"
+version = "9.2.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "decorator" },
+    { name = "ipython-pygments-lexers" },
+    { name = "jedi" },
+    { name = "matplotlib-inline" },
+    { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "prompt-toolkit" },
+    { name = "pygments" },
+    { name = "stack-data" },
+    { name = "traitlets" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/9d/02/63a84444a7409b3c0acd1de9ffe524660e0e5d82ee473e78b45e5bfb64a4/ipython-9.2.0.tar.gz", hash = "sha256:62a9373dbc12f28f9feaf4700d052195bf89806279fc8ca11f3f54017d04751b", size = 4424394, upload-time = "2025-04-25T17:55:40.498Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/78/ce/5e897ee51b7d26ab4e47e5105e7368d40ce6cfae2367acdf3165396d50be/ipython-9.2.0-py3-none-any.whl", hash = "sha256:fef5e33c4a1ae0759e0bba5917c9db4eb8c53fee917b6a526bd973e1ca5159f6", size = 604277, upload-time = "2025-04-25T17:55:37.625Z" },
+]
+
+[[package]]
+name = "ipython-pygments-lexers"
+version = "1.1.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" },
+]
+
+[[package]]
+name = "jedi"
+version = "0.19.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "parso" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/72/3a/79a912fbd4d8dd6fbb02bf69afd3bb72cf0c729bb3063c6f4498603db17a/jedi-0.19.2.tar.gz", hash = "sha256:4770dc3de41bde3966b02eb84fbcf557fb33cce26ad23da12c742fb50ecb11f0", size = 1231287, upload-time = "2024-11-11T01:41:42.873Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c0/5a/9cac0c82afec3d09ccd97c8b6502d48f165f9124db81b4bcb90b4af974ee/jedi-0.19.2-py2.py3-none-any.whl", hash = "sha256:a8ef22bde8490f57fe5c7681a3c83cb58874daf72b4784de3cce5b6ef6edb5b9", size = 1572278, upload-time = "2024-11-11T01:41:40.175Z" },
+]
+
 [[package]]
 name = "jinja2"
 version = "3.1.6"
@@ -730,6 +810,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c1/80/a61f99dc3a936413c3ee4e1eecac96c0da5ed07ad56fd975f1a9da5bc630/MarkupSafe-3.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:8e06879fc22a25ca47312fbe7c8264eb0b662f6db27cb2d3bbbc74b1df4b9b87", size = 15601, upload-time = "2024-10-18T15:21:23.499Z" },
 ]
 
+[[package]]
+name = "matplotlib-inline"
+version = "0.1.7"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "traitlets" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/99/5b/a36a337438a14116b16480db471ad061c36c3694df7c2084a0da7ba538b7/matplotlib_inline-0.1.7.tar.gz", hash = "sha256:8423b23ec666be3d16e16b60bdd8ac4e86e840ebd1dd11a30b9f117f2fa0ab90", size = 8159, upload-time = "2024-04-15T13:44:44.803Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8f/8e/9ad090d3553c280a8060fbf6e24dc1c0c29704ee7d1c372f0c174aa59285/matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca", size = 9899, upload-time = "2024-04-15T13:44:43.265Z" },
+]
+
 [[package]]
 name = "mbstrdecoder"
 version = "1.1.4"
@@ -1096,6 +1188,15 @@ wheels = [
     { url = "https://download.pytorch.org/whl/nightly/pandas-2.2.3-cp312-cp312-win_amd64.whl" },
 ]
 
+[[package]]
+name = "parso"
+version = "0.8.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/66/94/68e2e17afaa9169cf6412ab0f28623903be73d1b32e208d9e8e541bb086d/parso-0.8.4.tar.gz", hash = "sha256:eb3a7b58240fb99099a345571deecc0f9540ea5f4dd2fe14c2a99d6b281ab92d", size = 400609, upload-time = "2024-04-05T09:43:55.897Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c6/ac/dac4a63f978e4dcb3c6d3a78c4d8e0192a113d288502a1216950c41b1027/parso-0.8.4-py2.py3-none-any.whl", hash = "sha256:a418670a20291dacd2dddc80c377c5c3791378ee1e8d12bffc35420643d43f18", size = 103650, upload-time = "2024-04-05T09:43:53.299Z" },
+]
+
 [[package]]
 name = "pathvalidate"
 version = "3.2.3"
@@ -1126,6 +1227,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/68/85/8e6ea3d1089f2b6de3c1cd34bbbd7560912af9d34b057be3b8b8fefe1da3/peft-0.15.2-py3-none-any.whl", hash = "sha256:0dfc942b03b7af4b7267cd4e30b15e3a4a1d277adc581ce6245fc13f1f93d0a0", size = 411051, upload-time = "2025-04-15T15:27:50.799Z" },
 ]
 
+[[package]]
+name = "pexpect"
+version = "4.9.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "ptyprocess" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523", size = 63772, upload-time = "2023-11-25T06:56:14.81Z" },
+]
+
 [[package]]
 name = "platformdirs"
 version = "4.3.7"
@@ -1147,6 +1260,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f7/60/1974cfdd5bb770568ddc6f89f3e0df4cfdd1acffd5a609dff5e95f48c6e2/portalocker-3.1.1-py3-none-any.whl", hash = "sha256:80e984e24de292ff258a5bea0e4f3f778fff84c0ae1275dbaebc4658de4aacb3", size = 19661, upload-time = "2024-12-31T14:22:47.019Z" },
 ]
 
+[[package]]
+name = "prompt-toolkit"
+version = "3.0.51"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "wcwidth" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/bb/6e/9d084c929dfe9e3bfe0c6a47e31f78a25c54627d64a66e884a8bf5474f1c/prompt_toolkit-3.0.51.tar.gz", hash = "sha256:931a162e3b27fc90c86f1b48bb1fb2c528c2761475e57c9c06de13311c7b54ed", size = 428940, upload-time = "2025-04-15T09:18:47.731Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ce/4f/5249960887b1fbe561d9ff265496d170b55a735b76724f10ef19f9e40716/prompt_toolkit-3.0.51-py3-none-any.whl", hash = "sha256:52742911fde84e2d423e2f9a4cf1de7d7ac4e51958f648d9540e0fb8db077b07", size = 387810, upload-time = "2025-04-15T09:18:44.753Z" },
+]
+
 [[package]]
 name = "propcache"
 version = "0.3.1"
@@ -1201,6 +1326,40 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885, upload-time = "2025-02-13T21:54:37.486Z" },
 ]
 
+[[package]]
+name = "ptyprocess"
+version = "0.7.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" },
+]
+
+[[package]]
+name = "pudb"
+version = "2025.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "jedi" },
+    { name = "packaging" },
+    { name = "pygments" },
+    { name = "urwid" },
+    { name = "urwid-readline" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/93/40/8d17b16a1c2a36d8cc1befb5eda13310ad20fcb347e58bef5e104696e14a/pudb-2025.1.tar.gz", hash = "sha256:a528b29c69ce8b182a337872c5f046071f6d68d3415c6d7bf53bd27c264f58d0", size = 220623, upload-time = "2025-05-06T20:43:18.306Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b1/01/069766294390d3e10c77dfb553171466d67ffb51bf72a437650c0a5db86a/pudb-2025.1-py3-none-any.whl", hash = "sha256:f642d42e6054c992b43c463742650aa879fe290d7d7ffdeb21f7d00dc4587a21", size = 89208, upload-time = "2025-05-06T20:43:17.101Z" },
+]
+
+[[package]]
+name = "pure-eval"
+version = "0.2.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/cd/05/0a34433a064256a578f1783a10da6df098ceaa4a57bbeaa96a6c0352786b/pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42", size = 19752, upload-time = "2024-07-21T12:58:21.801Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" },
+]
+
 [[package]]
 name = "pyarrow"
 version = "20.0.0"
@@ -1660,6 +1819,20 @@ version = "2.1.0"
 source = { registry = "https://pypi.org/simple" }
 sdist = { url = "https://files.pythonhosted.org/packages/12/9a/7620d1e9dcb02839ed6d4b14064e609cdd7a8ae1e47289aa0456796dd9ca/sqlitedict-2.1.0.tar.gz", hash = "sha256:03d9cfb96d602996f1d4c2db2856f1224b96a9c431bdd16e78032a72940f9e8c", size = 21846, upload-time = "2022-12-03T13:39:13.102Z" }
 
+[[package]]
+name = "stack-data"
+version = "0.6.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "asttokens" },
+    { name = "executing" },
+    { name = "pure-eval" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/28/e3/55dcc2cfbc3ca9c29519eb6884dd1415ecb53b0e934862d3559ddcb7e20b/stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9", size = 44707, upload-time = "2023-09-30T13:58:05.479Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695", size = 24521, upload-time = "2023-09-30T13:58:03.53Z" },
+]
+
 [[package]]
 name = "submitit"
 version = "1.5.2"
@@ -1850,6 +2023,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/25/7e/0d889fc6c84e3df6b69aaafe893fc77f69b3d968ac9ce574d1c62c688050/tqdm_multiprocess-0.0.11-py3-none-any.whl", hash = "sha256:3ebdf03e7a675150fa0bbceaa9c3c64b8cb556e9ffafa4fe6c078e51820524aa", size = 9817, upload-time = "2020-10-27T06:57:53.167Z" },
 ]
 
+[[package]]
+name = "traitlets"
+version = "5.14.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/eb/79/72064e6a701c2183016abbbfedaba506d81e30e232a68c9f0d6f6fcd1574/traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7", size = 161621, upload-time = "2024-04-19T11:11:49.746Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" },
+]
+
 [[package]]
 name = "transformers"
 version = "4.51.3"
@@ -1944,6 +2126,27 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/6b/11/cc635220681e93a0183390e26485430ca2c7b5f9d33b15c74c2861cb8091/urllib3-2.4.0-py3-none-any.whl", hash = "sha256:4e16665048960a0900c702d4a66415956a584919c03361cac9f1df5c5dd7e813", size = 128680, upload-time = "2025-04-10T15:23:37.377Z" },
 ]
 
+[[package]]
+name = "urwid"
+version = "3.0.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "wcwidth" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/46/2d/71550379ed6b34968e14f73b0cf8574dee160acb6b820a066ab238ef2d4f/urwid-3.0.2.tar.gz", hash = "sha256:e7cb70ba1e7ff45779a5a57e43c57581ee7de6ceefb56c432491a4a6ce81eb78", size = 855353, upload-time = "2025-05-07T10:48:51.381Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b5/ee/2956918f14fd6e4310f7200108b53917f4da713d74c7ccd0d91a2e3a4f18/urwid-3.0.2-py3-none-any.whl", hash = "sha256:94ec1448d0178c881c01845c2b478cdc89f7b71bb65349466dbc99da1965eaac", size = 295994, upload-time = "2025-05-07T10:48:49.173Z" },
+]
+
+[[package]]
+name = "urwid-readline"
+version = "0.15.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "urwid" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ad/70/be318554495555eba7d8ff6e489f6f74ddb225b24086ba4af62a82e723fd/urwid_readline-0.15.1.tar.gz", hash = "sha256:9301444b86d58f7d26388506b704f142cefd193888488b4070d3a0fdfcfc0f84", size = 9007, upload-time = "2024-09-22T17:51:55.144Z" }
+
 [[package]]
 name = "viztracer"
 version = "1.0.3"
@@ -1994,6 +2197,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/00/46/46a2e7cf5113671eb7323dc1ae7a33b7e6ecbc1042b078f34f28a8968fa1/wandb-0.19.10-py3-none-win_amd64.whl", hash = "sha256:ae6b09efaa3c5b70146710fc153889f4faf059ea2e0bd2c1fb821dc59a77f50a", size = 20732825, upload-time = "2025-04-22T21:50:38.317Z" },
 ]
 
+[[package]]
+name = "wcwidth"
+version = "0.2.13"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/6c/63/53559446a878410fc5a5974feb13d31d78d752eb18aeba59c7fef1af7598/wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5", size = 101301, upload-time = "2024-01-06T02:10:57.829Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fd/84/fd2ba7aafacbad3c4201d395674fc6348826569da3c0937e75505ead3528/wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859", size = 34166, upload-time = "2024-01-06T02:10:55.763Z" },
+]
+
 [[package]]
 name = "win32-setctime"
 version = "1.2.0"