Move components to torchbenchmark/_components (#2201)

Summary:
When running as a package and user changed to a different directory, we need to also package the `components` directory.

Fixes #2049

Pull Request resolved: #2201

Reviewed By: aaronenyeshi

Differential Revision: D55024021

Pulled By: xuzhao9

fbshipit-source-id: 1352b075b376b14a885e2eb01cb72e627958f17c

run.py

@@ -175,7 +175,7 @@ def run_one_step(
     result_summary = []
     flops_model_analyzer = None
     if "flops" in metrics_needed:
-        from components.model_analyzer.TorchBenchAnalyzer import ModelAnalyzer
+        from torchbenchmark._components.model_analyzer.TorchBenchAnalyzer import ModelAnalyzer
 
         flops_model_analyzer = ModelAnalyzer(
             export_metrics_file, ["flops"], metrics_gpu_backend
@@ -554,11 +554,11 @@ def main() -> None:
     metrics_gpu_backend = args.metrics_gpu_backend
     if metrics_needed:
         if metrics_gpu_backend == "dcgm":
-            from components.model_analyzer.TorchBenchAnalyzer import check_dcgm
+            from torchbenchmark._components.model_analyzer.TorchBenchAnalyzer import check_dcgm
 
             check_dcgm()
         elif "gpu_peak_mem" in metrics_needed:
-            from components.model_analyzer.TorchBenchAnalyzer import check_nvml
+            from torchbenchmark._components.model_analyzer.TorchBenchAnalyzer import check_nvml
 
             check_nvml()
         if "gpu_peak_mem" in metrics_needed or (

scripts/torchbench_test.sh

@@ -17,8 +17,8 @@ parent_dir=$(dirname "$(readlink -f "$0")")/..
 cd ${parent_dir}
 
 # Test subprocess worker
-python -m components.test.test_subprocess
-python -m components.test.test_worker
+python -m torchbenchmark._components.test.test_subprocess
+python -m torchbenchmark._components.test.test_worker
 
 # Test models
 python test.py -v

test_bench.py

@@ -15,7 +15,7 @@
 import pytest
 import time
 import torch
-from components._impl.workers import subprocess_worker
+from torchbenchmark._components._impl.workers import subprocess_worker
 from torchbenchmark import _list_model_paths, ModelTask, get_metadata_from_yaml
 from torchbenchmark.util.machine_config import get_machine_state
 from torchbenchmark.util.metadata_utils import skip_by_metadata

test_imports.py

@@ -1,4 +1,4 @@
-import torchbenchmark.models.densenet121
+import torchbenchmark.models
 
 model, example_inputs = torchbenchmark.models.densenet121.Model(
     test="eval", device="cuda", batch_size=1

torchbenchmark/__init__.py

@@ -14,8 +14,8 @@
 
 import torch
 
-from components._impl.tasks import base as base_task
-from components._impl.workers import subprocess_worker
+from ._components._impl.tasks import base as base_task
+from ._components._impl.workers import subprocess_worker
 from . import models
 from . import canary_models
 from . import e2e_models

components/_impl/tasks/base.py → torchbenchmark/_components/_impl/tasks/base.py

@@ -7,7 +7,7 @@
 import textwrap
 import typing
 
-from components._impl.workers import base
+from torchbenchmark._components._impl.workers import base
 
 
 class TaskBase(abc.ABC):

components/_impl/workers/in_process_worker.py → torchbenchmark/_components/_impl/workers/in_process_worker.py

@@ -2,7 +2,7 @@
 import textwrap
 import typing
 
-from components._impl.workers import base
+from torchbenchmark._components._impl.workers import base
 
 
 class InProcessWorker(base.WorkerBase):

components/_impl/workers/subprocess_worker.py → torchbenchmark/_components/_impl/workers/subprocess_worker.py

@@ -12,8 +12,8 @@
 from pathlib import Path
 from typing import Optional, Dict, List, Any, Tuple
 
-from components._impl.workers import base
-from components._impl.workers import subprocess_rpc
+from torchbenchmark._components._impl.workers import base
+from torchbenchmark._components._impl.workers import subprocess_rpc
 
 
 class SubprocessWorker(base.WorkerBase):
@@ -211,7 +211,7 @@ def _bootstrap_worker(self) -> None:
                 # The parent gets priority, but a subclass could set PYTHONPATH
                 # so we have to respect extra paths.
                 sys.path.extend([i for i in sys_path_old if i and i not in sys.path])
-                from components._impl.workers import subprocess_rpc
+                from torchbenchmark._components._impl.workers import subprocess_rpc
                 output_pipe = subprocess_rpc.Pipe(
                     write_handle={self._output_pipe.write_handle})
                 output_pipe.write(subprocess_rpc.BOOTSTRAP_IMPORT_SUCCESS)

components/test/test_subprocess.py → torchbenchmark/_components/test/test_subprocess.py

@@ -14,8 +14,8 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 
 try:
-    from components._impl.tasks import base as task_base
-    from components._impl.workers import subprocess_rpc
+    from torchbenchmark._components._impl.tasks import base as task_base
+    from torchbenchmark._components._impl.workers import subprocess_rpc
 except (ImportError, ModuleNotFoundError):
     print(f"""
         This test must be run from the repo root directory as

components/test/test_worker.py → torchbenchmark/_components/test/test_worker.py

@@ -8,10 +8,10 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 
 try:
-    from components._impl.workers import base as base_worker
-    from components._impl.workers import in_process_worker
-    from components._impl.workers import subprocess_worker
-    from components._impl.workers import subprocess_rpc
+    from torchbenchmark._components._impl.workers import base as base_worker
+    from torchbenchmark._components._impl.workers import in_process_worker
+    from torchbenchmark._components._impl.workers import subprocess_worker
+    from torchbenchmark._components._impl.workers import subprocess_rpc
 except (ImportError, ModuleNotFoundError):
     print(
         f"""
@@ -232,7 +232,7 @@ def test_subprocess_worker(self) -> None:
         self._test_child_trace_exception(worker)
 
     def test_subprocess_worker_segv_handling(self):
-        worker = subprocess_worker.SubprocessWorker(timeout=1)
+        worker = subprocess_worker.SubprocessWorker(timeout=2)
         with self.assertRaisesRegex(
             OSError, f"Subprocess terminates with code {int(signal.SIGSEGV)}"
         ):

torchbenchmark/util/experiment/metrics.py

@@ -51,7 +51,7 @@ def get_latencies(func, device: str, nwarmup=WARMUP_ROUNDS, num_iter=BENCHMARK_I
 
 def get_peak_memory(func, device: str, num_iter=MEMPROF_ITER, export_metrics_file='', metrics_needed=[], metrics_gpu_backend='dcgm', cpu_monitored_pid=None) -> Tuple[Optional[float], Optional[str], Optional[float]]:
     "Run one step of the model, and return the peak memory in MB."
-    from components.model_analyzer.TorchBenchAnalyzer import ModelAnalyzer
+    from torchbenchmark._components.model_analyzer.TorchBenchAnalyzer import ModelAnalyzer
     new_metrics_needed = [_ for _ in metrics_needed if _ in ['cpu_peak_mem', 'gpu_peak_mem']]
     if not new_metrics_needed:
         raise ValueError(f"Expected metrics_needed to be non-empty, get: {metrics_needed}")

userbenchmark/model-stableness/run.py

@@ -10,7 +10,7 @@
 from ..utils import REPO_PATH, add_path, get_output_dir, get_output_json, dump_output
 
 with add_path(REPO_PATH):
-    from components._impl.workers.subprocess_rpc import UnserializableException, ChildTraceException
+    from torchbenchmark._components._impl.workers.subprocess_rpc import UnserializableException, ChildTraceException
     from torchbenchmark.util.experiment.instantiator import list_models, load_model_isolated, TorchBenchModelConfig
     from torchbenchmark.util.experiment.metrics import TorchBenchModelMetrics, get_model_test_metrics