diff --git a/clearml/__init__.py b/clearml/__init__.py index d127f0a5..29dde4ac 100644 --- a/clearml/__init__.py +++ b/clearml/__init__.py @@ -1,4 +1,5 @@ """ ClearML open SDK """ +from six import PY2 from .version import __version__ from .task import Task @@ -7,7 +8,31 @@ from .logger import Logger from .storage import StorageManager from .errors import UsageError from .datasets import Dataset -from .automation.controller import PipelineController -__all__ = ["__version__", "Task", "InputModel", "OutputModel", "Model", "Logger", - "StorageManager", "UsageError", "Dataset", "PipelineController"] +if not PY2: + from .automation.controller import PipelineController + + __all__ = [ + "__version__", + "Task", + "InputModel", + "OutputModel", + "Model", + "Logger", + "StorageManager", + "UsageError", + "Dataset", + "PipelineController", + ] +else: + __all__ = [ + "__version__", + "Task", + "InputModel", + "OutputModel", + "Model", + "Logger", + "StorageManager", + "UsageError", + "Dataset", + ]