添加ifbench对nltk包的判断

ais_bench/benchmark/configs/datasets/ifbench/ifbench_0_shot_gen_str.py

@@ -26,6 +26,7 @@
         abbr='ifbench',
         type=IFBenchDataset,
         path='ais_bench/datasets/ifbench/data/train-00000-of-00001.parquet',
+        nltk_path='/path/to/nltk_data',
         reader_cfg=ifbench_reader_cfg,
         infer_cfg=ifbench_infer_cfg,
         eval_cfg=ifbench_eval_cfg,

ais_bench/benchmark/datasets/ifbench/instructions_util.py

@@ -5,18 +5,33 @@
 import re
 
 
+_NLTK_RESOURCE_PATHS = {
+    'punkt_tab': 'tokenizers/punkt_tab',
+    'averaged_perceptron_tagger_eng':
+        'taggers/averaged_perceptron_tagger_eng',
+    'stopwords': 'corpora/stopwords',
+}
+
+
 def _ensure_nltk_data(data_name):
-    """Ensure NLTK data is available, downloading if necessary."""
+    """Ensure the requested NLTK resource is available."""
+    resource_path = _NLTK_RESOURCE_PATHS.get(
+        data_name, f'tokenizers/{data_name}'
+    )
+
     try:
-        nltk.data.find(f'tokenizers/{data_name}')
+        nltk.data.find(resource_path)
     except (LookupError, OSError):
-        try:
-            nltk.download(data_name, quiet=True)
-        except (LookupError, OSError):
-            # Fall back to downloading the full 'punkt' package if the
-            # specific resource (e.g. 'punkt_tab') is not available as a
-            # standalone download in the installed NLTK version.
-            nltk.download('punkt', quiet=True)
+        # 在线环境下仍可自动下载；离线环境预装正确后不会进入这里
+        downloaded = nltk.download(data_name, quiet=True)
+        if not downloaded:
+            raise LookupError(
+                f"Missing NLTK resource: {resource_path}. "
+                f"Please install it under an NLTK_DATA directory."
+            )
+
+        # 下载后再检查，避免静默失败
+        nltk.data.find(resource_path)
 
 
 WORD_LIST = [
@@ -248,8 +263,6 @@ def _ensure_nltk_data(data_name):
     'injury', 'insect', 'surprise', 'apartment',
 ]  # pylint: disable=line-too-long
 
-_ensure_nltk_data('punkt_tab')
-
 _ALPHABETS = '([A-Za-z])'
 _PREFIXES = '(Mr|St|Mrs|Ms|Dr)[.]'
 _SUFFIXES = '(Inc|Ltd|Jr|Sr|Co)'

ais_bench/benchmark/runners/local.py

@@ -66,6 +66,46 @@ def __init__(self,
         for k, v in kwargs.items():
             self.logger.warning(f'Ignored argument in {self.__module__}: {k}={v}')
 
+    def _get_subprocess_env(self, task) -> Dict[str, str]:
+        env = os.environ.copy()
+        configured_paths = []
+        for dataset_cfg in getattr(task, 'dataset_cfgs', []):
+            if 'nltk_path' not in dataset_cfg:
+                continue
+            raw_path = dataset_cfg['nltk_path']
+            if not isinstance(raw_path, str) or not raw_path.strip():
+                message = (
+                    f"Task {task.name} has invalid nltk_path: {raw_path!r}; "
+                    "expected a non-empty string")
+                self.logger.error(message)
+                raise ParameterValueError(
+                    RUNNER_CODES.INVALID_NLTK_PATH, message)
+            path = osp.abspath(osp.expandvars(osp.expanduser(raw_path.strip())))
+            configured_paths.append(path)
+
+        distinct_paths = {osp.normcase(path) for path in configured_paths}
+        if len(distinct_paths) > 1:
+            message = (
+                f"Task {task.name} has conflicting nltk_path values: "
+                f"{configured_paths}")
+            self.logger.error(message)
+            raise ParameterValueError(RUNNER_CODES.INVALID_NLTK_PATH, message)
+        if not configured_paths:
+            return env
+
+        nltk_path = configured_paths[0]
+        if not osp.isdir(nltk_path) or not os.access(nltk_path, os.R_OK):
+            message = (
+                f"Task {task.name} has unusable nltk_path: {nltk_path}; "
+                "the path must be an existing readable directory")
+            self.logger.error(message)
+            raise ParameterValueError(RUNNER_CODES.INVALID_NLTK_PATH, message)
+
+        env['NLTK_DATA'] = nltk_path
+        self.logger.info(
+            f"Task {task.name} sets NLTK_DATA={nltk_path} for its subprocess")
+        return env
+
     def launch(self, tasks: List[Dict[str, Any]]) -> List[Tuple[str, int]]:
         """Launch multiple tasks.
 
@@ -151,7 +191,8 @@ def _run_debug(self, tasks: List[Dict[str, Any]], all_gpu_ids: List[int], monito
                 tmpl = get_command_template(all_gpu_ids[:num_gpus])
                 cmd = task.get_command(cfg_path=param_file, template=tmpl)
 
-                proc = subprocess.Popen(cmd, shell=True, text=True)
+                env = self._get_subprocess_env(task)
+                proc = subprocess.Popen(cmd, shell=True, text=True, env=env)
                 try:
                     proc.wait()
                 except KeyboardInterrupt:
@@ -253,12 +294,14 @@ def _launch(self, task, gpu_ids, index):
             # Run command
             out_path = task.get_log_path(file_extension='out')
             mmengine.mkdir_or_exist(osp.split(out_path)[0])
+            env = self._get_subprocess_env(task)
             with open(out_path, 'w', encoding='utf-8') as stdout:
                 result = subprocess.run(cmd,
                                         shell=True,
                                         text=True,
                                         stdout=stdout,
-                                        stderr=stdout)
+                                        stderr=stdout,
+                                        env=env)
             if result.returncode != 0:
                 self.logger.error(RUNNER_CODES.TASK_FAILED, f"{task_name} failed with code {result.returncode}, see\n{out_path}")
         finally:

ais_bench/benchmark/utils/logging/error_codes.py

@@ -116,6 +116,9 @@ class SUMM_CODES:
 class RUNNER_CODES:
     UNKNOWN_ERROR = BaseErrorCode("RUNNER-UNK-001", ErrorModule.RUNNER, ErrorType.UNKNOWN, 1, "unknown error of runner")
     TASK_FAILED = BaseErrorCode("RUNNER-TASK-001", ErrorModule.RUNNER, ErrorType.TASK, 1, "task failed")  # docs coverd
+    INVALID_NLTK_PATH = BaseErrorCode(
+        'RUNNER-PARAM-001', ErrorModule.RUNNER, ErrorType.PARAM, 1,
+        'invalid NLTK data path')
 
 
 class TMON_CODES:

tests/UT/runners/test_local.py

@@ -6,6 +6,14 @@
 from mmengine.config import ConfigDict
 
 from ais_bench.benchmark.runners.local import LocalRunner, get_command_template
+from ais_bench.benchmark.utils.logging.exceptions import ParameterValueError
+
+
+def _task_with_datasets(*datasets):
+    task = MagicMock()
+    task.name = 'test_task'
+    task.dataset_cfgs = [ConfigDict(dataset) for dataset in datasets]
+    return task
 
 
 class TestGetCommandTemplate(unittest.TestCase):
@@ -64,6 +72,104 @@ def setUp(self):
         self.max_num_workers = 4
         self.max_workers_per_gpu = 1
 
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    def test_get_subprocess_env_sets_nltk_data_without_mutating_parent(
+            self, mock_logger_class):
+        """A valid dataset path is private to the child environment."""
+        runner = LocalRunner(task=self.task_cfg)
+        with tempfile.TemporaryDirectory() as nltk_dir:
+            parent_before = os.environ.get('NLTK_DATA')
+            env = runner._get_subprocess_env(
+                _task_with_datasets({'nltk_path': nltk_dir}))
+
+        self.assertEqual(env['NLTK_DATA'], os.path.abspath(nltk_dir))
+        self.assertEqual(os.environ.get('NLTK_DATA'), parent_before)
+        mock_logger_class.return_value.info.assert_called_once()
+        mock_logger_class.return_value.error.assert_not_called()
+
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    def test_get_subprocess_env_preserves_inherited_value_when_unconfigured(
+            self, mock_logger_class):
+        """Tasks without the setting retain the parent process value."""
+        runner = LocalRunner(task=self.task_cfg)
+        with patch.dict(os.environ, {'NLTK_DATA': 'inherited-value'}):
+            env = runner._get_subprocess_env(_task_with_datasets({}))
+            self.assertEqual(os.environ.get('NLTK_DATA'), 'inherited-value')
+
+        self.assertEqual(env['NLTK_DATA'], 'inherited-value')
+        mock_logger_class.return_value.info.assert_not_called()
+        mock_logger_class.return_value.error.assert_not_called()
+
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    def test_get_subprocess_env_rejects_conflicting_paths(
+            self, mock_logger_class):
+        """Distinct dataset paths cannot be selected for one child task."""
+        runner = LocalRunner(task=self.task_cfg)
+        with tempfile.TemporaryDirectory() as first, \
+                tempfile.TemporaryDirectory() as second:
+            with self.assertRaises(ParameterValueError) as context:
+                runner._get_subprocess_env(_task_with_datasets(
+                    {'nltk_path': first}, {'nltk_path': second}))
+
+        self.assertEqual(context.exception.error_code_str, 'RUNNER-PARAM-001')
+        mock_logger_class.return_value.error.assert_called_once()
+        mock_logger_class.return_value.info.assert_not_called()
+
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    def test_get_subprocess_env_expands_environment_path(self,
+                                                         mock_logger_class):
+        """Configured paths support environment-variable expansion."""
+        runner = LocalRunner(task=self.task_cfg)
+        with tempfile.TemporaryDirectory() as nltk_dir:
+            with patch.dict(
+                    os.environ,
+                    {'AISBENCH_NLTK_TEST_PATH': nltk_dir}):
+                env = runner._get_subprocess_env(_task_with_datasets(
+                    {'nltk_path': '%AISBENCH_NLTK_TEST_PATH%'}))
+
+        self.assertEqual(env['NLTK_DATA'], os.path.abspath(nltk_dir))
+        mock_logger_class.return_value.info.assert_called_once()
+        mock_logger_class.return_value.error.assert_not_called()
+
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    def test_get_subprocess_env_rejects_invalid_configured_path(
+            self, mock_logger_class):
+        """Malformed and unusable configured paths are rejected and logged."""
+        runner = LocalRunner(task=self.task_cfg)
+        missing = os.path.join(
+            tempfile.gettempdir(), 'aisbench-missing-nltk-data')
+        invalid_values = [None, '', 123, missing]
+
+        for value in invalid_values:
+            with self.subTest(value=value):
+                with self.assertRaises(ParameterValueError) as context:
+                    runner._get_subprocess_env(
+                        _task_with_datasets({'nltk_path': value}))
+                self.assertEqual(
+                    context.exception.error_code_str, 'RUNNER-PARAM-001')
+
+        with tempfile.NamedTemporaryFile() as regular_file:
+            with self.assertRaises(ParameterValueError) as context:
+                runner._get_subprocess_env(
+                    _task_with_datasets({'nltk_path': regular_file.name}))
+            self.assertEqual(context.exception.error_code_str, 'RUNNER-PARAM-001')
+
+        with tempfile.TemporaryDirectory() as unreadable:
+            with patch(
+                    'ais_bench.benchmark.runners.local.os.access',
+                    return_value=False):
+                with self.assertRaises(ParameterValueError) as context:
+                    runner._get_subprocess_env(
+                        _task_with_datasets({'nltk_path': unreadable}))
+                self.assertEqual(
+                    context.exception.error_code_str, 'RUNNER-PARAM-001')
+
+        self.assertEqual(
+            mock_logger_class.return_value.error.call_count,
+            len(invalid_values) + 2,
+        )
+        mock_logger_class.return_value.info.assert_not_called()
+
     @patch('ais_bench.benchmark.runners.base.AISLogger')
     def test_init(self, mock_logger_class):
         """测试LocalRunner初始化"""
@@ -212,6 +318,36 @@ def test_launch_with_visible_devices(self, mock_environ, mock_logger_class):
                             call_args = str(mock_logger.debug.call_args_list)
                             self.assertIn("Available devices", call_args)
 
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    @patch('ais_bench.benchmark.runners.local.TASKS')
+    def test_run_debug_passes_nltk_environment(self, mock_tasks,
+                                               mock_logger_class):
+        """Debug subprocess receives the configured NLTK data path."""
+        runner = LocalRunner(task=self.task_cfg, debug=True)
+        mock_task = MagicMock()
+        mock_task.name = 'test_task'
+        mock_task.num_gpus = 0
+        mock_task.get_command.return_value = 'python test.py'
+        mock_task.cfg.dump = MagicMock()
+        mock_tasks.build.return_value = mock_task
+
+        with tempfile.TemporaryDirectory() as nltk_dir:
+            mock_task.dataset_cfgs = [ConfigDict({'nltk_path': nltk_dir})]
+            with patch(
+                    'ais_bench.benchmark.runners.local.subprocess.Popen'
+            ) as subprocess_mock, patch(
+                    'ais_bench.benchmark.runners.local.os.remove'
+            ), patch(
+                    'ais_bench.benchmark.runners.local.mmengine.mkdir_or_exist'
+            ), patch('uuid.uuid4', return_value=MagicMock(hex='test')):
+                subprocess_mock.return_value.wait.return_value = None
+                runner._run_debug([{'work_dir': '/tmp/test', 'cli_args': {}}],
+                                  [], MagicMock())
+
+            self.assertEqual(
+                subprocess_mock.call_args.kwargs['env']['NLTK_DATA'],
+                os.path.abspath(nltk_dir))
+
     @patch('ais_bench.benchmark.runners.base.AISLogger')
     @patch('ais_bench.benchmark.runners.local.TASKS')
     def test_run_debug(self, mock_tasks, mock_logger_class):
@@ -257,6 +393,33 @@ def test_run_debug(self, mock_tasks, mock_logger_class):
                             self.assertEqual(status[0][0], "test_task")
                             self.assertEqual(status[0][1], 0)
 
+    @patch('ais_bench.benchmark.runners.base.AISLogger')
+    def test_launch_passes_nltk_environment(self, mock_logger_class):
+        """Normal subprocess receives the configured NLTK data path."""
+        runner = LocalRunner(task=self.task_cfg, debug=False)
+        mock_task = MagicMock()
+        mock_task.name = 'test_task'
+        mock_task.get_command.return_value = 'python test.py'
+        mock_task.get_log_path.return_value = '/tmp/test.out'
+        mock_task.cfg.dump = MagicMock()
+
+        with tempfile.TemporaryDirectory() as nltk_dir:
+            mock_task.dataset_cfgs = [ConfigDict({'nltk_path': nltk_dir})]
+            with patch(
+                    'ais_bench.benchmark.runners.local.subprocess.run'
+            ) as subprocess_mock, patch(
+                    'ais_bench.benchmark.runners.local.os.remove'
+            ), patch(
+                    'ais_bench.benchmark.runners.local.mmengine.mkdir_or_exist'
+            ), patch('uuid.uuid4', return_value=MagicMock(hex='test')), patch(
+                    'builtins.open', create=True):
+                subprocess_mock.return_value.returncode = 0
+                runner._launch(mock_task, [0], 0)
+
+            self.assertEqual(
+                subprocess_mock.call_args.kwargs['env']['NLTK_DATA'],
+                os.path.abspath(nltk_dir))
+
     @patch('ais_bench.benchmark.runners.base.AISLogger')
     def test_launch_method(self, mock_logger_class):
         """测试_launch方法"""