[Setup] Fix dependency (#39)

setup.py

@@ -24,13 +24,16 @@ def setup():
 
     setuptools.setup(
         name="slapo",
-        description="Slapo: A Scahedule LAnguage for Progressive Optimization.",
+        description="Slapo: A Schedule LAnguage for Progressive Optimization.",
         version=get_version(),
         author="Slapo Community",
         long_description=long_description,
         long_description_content_type="text/markdown",
         setup_requires=[],
-        install_requires=[],
+        install_requires=[
+            "packaging",
+            "psutil",
+        ],
         packages=setuptools.find_packages(),
         url="https://github.com/awslabs/slapo",
         python_requires=">=3.7",

slapo/model_dialect/deepspeed/engine.py

@@ -2,11 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from deepspeed.runtime.pipe.topology import (
-    PipeModelDataParallelTopology,
-    PipelineParallelGrid,
-)
-
 from ..registry import register_model_dialect
 from ...logger import get_logger, INFO
 
@@ -17,6 +12,10 @@
 def init_ds_engine(model, **kwargs):
     """Initialize the DeepSpeed engine."""
     import deepspeed
+    from deepspeed.runtime.pipe.topology import (
+        PipeModelDataParallelTopology,
+        PipelineParallelGrid,
+    )
 
     if "config" not in kwargs:
         raise ValueError("DeepSpeed config not provided.")

slapo/tracer.py

@@ -15,12 +15,6 @@
 from torch.fx._symbolic_trace import HAS_VARSTUFF, PH, _assert_is_none, _patch_function
 from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
 from torch.fx.node import base_types
-from transformers.utils.fx import (
-    _IS_IN_DEBUG_MODE,
-    _MANUAL_META_OVERRIDES,
-    Proxy,
-    _proxies_to_metas,
-)
 
 from .logger import get_logger
 
@@ -281,7 +275,13 @@ def trace(model: nn.Module, **kwargs: dict[str, Any]):
     logger.debug("Tracer: %s Model: %s", tracer_cls_name, model.__class__.__name__)
     if isinstance(tracer_cls_name, str):
         if tracer_cls_name == "huggingface":
-            from transformers.utils.fx import HFTracer
+            from transformers.utils.fx import (
+                HFTracer,
+                _IS_IN_DEBUG_MODE,
+                _MANUAL_META_OVERRIDES,
+                Proxy,
+                _proxies_to_metas,
+            )
 
             assert (
                 "concrete_args" in kwargs