Create new distilabel.constants module to store constants and avoid…

… circular imports (#861)

* Bump version to `1.4.0`

* Refactor constants modules to a higher level to avoid circular imports

* Update src/distilabel/__init__.py

Co-authored-by: Gabriel Martín Blázquez <[email protected]>

---------

Co-authored-by: Gabriel Martín Blázquez <[email protected]>

src/distilabel/__init__.py

@@ -14,6 +14,6 @@
 
 from rich import traceback as rich_traceback
 
-__version__ = "1.3.0"
+__version__ = "1.3.1"
 
 rich_traceback.install(show_locals=True)

src/distilabel/cli/pipeline/utils.py

@@ -23,10 +23,7 @@
 from pydantic import HttpUrl, ValidationError
 from pydantic.type_adapter import TypeAdapter
 
-from distilabel.pipeline.constants import (
-    ROUTING_BATCH_FUNCTION_ATTR_NAME,
-    STEP_ATTR_NAME,
-)
+from distilabel.constants import ROUTING_BATCH_FUNCTION_ATTR_NAME, STEP_ATTR_NAME
 from distilabel.pipeline.local import Pipeline
 
 if TYPE_CHECKING:

src/distilabel/constants.py

@@ -0,0 +1,37 @@
+# Copyright 2023-present, Argilla, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Final
+
+# Steps related constants
+DISTILABEL_METADATA_KEY: Final[str] = "distilabel_metadata"
+
+# Pipeline related constants
+STEP_ATTR_NAME: Final[str] = "step"
+INPUT_QUEUE_ATTR_NAME: Final[str] = "input_queue"
+RECEIVES_ROUTED_BATCHES_ATTR_NAME: Final[str] = "receives_routed_batches"
+ROUTING_BATCH_FUNCTION_ATTR_NAME: Final[str] = "routing_batch_function"
+CONVERGENCE_STEP_ATTR_NAME: Final[str] = "convergence_step"
+LAST_BATCH_SENT_FLAG: Final[str] = "last_batch_sent"
+
+
+__all__ = [
+    "STEP_ATTR_NAME",
+    "INPUT_QUEUE_ATTR_NAME",
+    "RECEIVES_ROUTED_BATCHES_ATTR_NAME",
+    "ROUTING_BATCH_FUNCTION_ATTR_NAME",
+    "CONVERGENCE_STEP_ATTR_NAME",
+    "LAST_BATCH_SENT_FLAG",
+    "DISTILABEL_METADATA_KEY",
+]

src/distilabel/distiset.py

@@ -29,7 +29,7 @@
 from pyarrow.lib import ArrowInvalid
 from typing_extensions import Self
 
-from distilabel.pipeline.constants import STEP_ATTR_NAME
+from distilabel.constants import STEP_ATTR_NAME
 from distilabel.utils.card.dataset_card import (
     DistilabelDatasetCard,
     size_categories_parser,
@@ -536,7 +536,7 @@ def create_distiset(  # noqa: C901
         >>> distiset = create_distiset(Path.home() / ".cache/distilabel/pipelines/path-to-pipe-hashname")
         ```
     """
-    from distilabel.steps.constants import DISTILABEL_METADATA_KEY
+    from distilabel.constants import DISTILABEL_METADATA_KEY
 
     logger = logging.getLogger("distilabel.distiset")
 

src/distilabel/pipeline/_dag.py

@@ -30,7 +30,7 @@
 
 import networkx as nx
 
-from distilabel.pipeline.constants import (
+from distilabel.constants import (
     CONVERGENCE_STEP_ATTR_NAME,
     ROUTING_BATCH_FUNCTION_ATTR_NAME,
     STEP_ATTR_NAME,

src/distilabel/pipeline/base.py

@@ -39,19 +39,19 @@
 from upath import UPath
 
 from distilabel import __version__
-from distilabel.distiset import create_distiset
-from distilabel.mixins.requirements import RequirementsMixin
-from distilabel.pipeline._dag import DAG
-from distilabel.pipeline.batch import _Batch
-from distilabel.pipeline.batch_manager import _BatchManager
-from distilabel.pipeline.constants import (
+from distilabel.constants import (
     CONVERGENCE_STEP_ATTR_NAME,
     INPUT_QUEUE_ATTR_NAME,
     LAST_BATCH_SENT_FLAG,
     RECEIVES_ROUTED_BATCHES_ATTR_NAME,
     ROUTING_BATCH_FUNCTION_ATTR_NAME,
     STEP_ATTR_NAME,
 )
+from distilabel.distiset import create_distiset
+from distilabel.mixins.requirements import RequirementsMixin
+from distilabel.pipeline._dag import DAG
+from distilabel.pipeline.batch import _Batch
+from distilabel.pipeline.batch_manager import _BatchManager
 from distilabel.pipeline.write_buffer import _WriteBuffer
 from distilabel.steps.base import GeneratorStep
 from distilabel.steps.generators.utils import make_generator_step

src/distilabel/pipeline/batch_manager.py

@@ -17,12 +17,12 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple, Union
 
-from distilabel.pipeline._dag import DAG
-from distilabel.pipeline.batch import _Batch
-from distilabel.pipeline.constants import (
+from distilabel.constants import (
     RECEIVES_ROUTED_BATCHES_ATTR_NAME,
     STEP_ATTR_NAME,
 )
+from distilabel.pipeline._dag import DAG
+from distilabel.pipeline.batch import _Batch
 from distilabel.steps.base import _Step
 from distilabel.utils.files import list_files_in_dir
 from distilabel.utils.serialization import (

src/distilabel/pipeline/ray.py

@@ -15,10 +15,10 @@
 import sys
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
+from distilabel.constants import INPUT_QUEUE_ATTR_NAME
 from distilabel.distiset import create_distiset
 from distilabel.llms.vllm import vLLM
 from distilabel.pipeline.base import BasePipeline
-from distilabel.pipeline.constants import INPUT_QUEUE_ATTR_NAME
 from distilabel.pipeline.step_wrapper import _StepWrapper
 from distilabel.utils.logging import setup_logging, stop_logging
 from distilabel.utils.serialization import TYPE_INFO_KEY

src/distilabel/pipeline/step_wrapper.py

@@ -16,9 +16,9 @@
 from queue import Queue
 from typing import Any, Dict, List, Optional, Union, cast
 
+from distilabel.constants import LAST_BATCH_SENT_FLAG
 from distilabel.llms.mixins.cuda_device_placement import CudaDevicePlacementMixin
 from distilabel.pipeline.batch import _Batch
-from distilabel.pipeline.constants import LAST_BATCH_SENT_FLAG
 from distilabel.pipeline.typing import StepLoadStatus
 from distilabel.steps.base import GeneratorStep, Step, _Step
 

src/distilabel/steps/tasks/base.py

@@ -18,6 +18,7 @@
 from pydantic import Field
 from typing_extensions import override
 
+from distilabel.constants import DISTILABEL_METADATA_KEY
 from distilabel.llms.base import LLM
 from distilabel.mixins.runtime_parameters import RuntimeParameter
 from distilabel.steps.base import (
@@ -26,7 +27,6 @@
     StepInput,
     _Step,
 )
-from distilabel.steps.constants import DISTILABEL_METADATA_KEY
 from distilabel.utils.dicts import group_dicts
 
 if TYPE_CHECKING: