uber-compose 2.3.3__tar.gz → 3.0.0a1__tar.gz

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: uber-compose
-Version: 2.3.3
+Version: 3.0.0a1
 Summary: docker compose testing env orchestrator
 Home-page: https://github.com/ko10ok/uber-compose
 Author: Yuriy Sagitov
@@ -231,6 +231,7 @@ This is particularly useful when:
 - **[E2E Test Setup](docs/E2E_CONTAINER_SETUP.md)** - Configure test container for E2E testing
 - **[CLI Usage Guide](docs/CLI_USAGE.md)** - Guide for using CommonJsonCli with JSON log parsing
 - **[X-Migration](docs/X_MIGRATION.md)** - Extended service initialization and migration commands
+- **[Lifecycle Hooks](docs/HOOKS.md)** - Register callbacks that run after the environment is up
 
 ## 🛠️ Development Guide
 

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/README.md

@@ -198,6 +198,7 @@ This is particularly useful when:
 - **[E2E Test Setup](docs/E2E_CONTAINER_SETUP.md)** - Configure test container for E2E testing
 - **[CLI Usage Guide](docs/CLI_USAGE.md)** - Guide for using CommonJsonCli with JSON log parsing
 - **[X-Migration](docs/X_MIGRATION.md)** - Extended service initialization and migration commands
+- **[Lifecycle Hooks](docs/HOOKS.md)** - Register callbacks that run after the environment is up
 
 ## 🛠️ Development Guide
 

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/__init__.py

@@ -10,6 +10,23 @@ from uber_compose.env_description.env_types import OverridenService
 from uber_compose.env_description.env_types import Service
 from uber_compose.helpers.exec_result import ExecTimeout
 from uber_compose.helpers.health_policy import UpHealthPolicy
+from uber_compose.helpers.hooks import UberComposeEvent
+from uber_compose.parallel_plugin import ExecutionPlan
+from uber_compose.parallel_plugin import VedroUberComposeParallel
+from uber_compose.parallel_plugin import WorkerAssignment
+from uber_compose.parallel_plugin import build_execution_plan
+from uber_compose.parallel_plugin import build_execution_plan_from_scheduler
+from uber_compose.parallel_plugin import env_group_key
+from uber_compose.parallel_plugin import env_group_key_for
+from uber_compose.parallel_plugin import get_execution_plan
+from uber_compose.parallel_plugin import get_namespace_map
+from uber_compose.parallel_plugin import get_worker_groups
+from uber_compose.parallel_plugin import get_worker_id
+from uber_compose.parallel_plugin import get_workers_total
+from uber_compose.parallel_plugin import is_parallel_worker
+from uber_compose.parallel_plugin import is_parallel_worker_process
+from uber_compose.parallel_plugin import reset_execution_plan
+from uber_compose.parallel_plugin import set_execution_plan
 from uber_compose.uber_compose import SystemUberCompose
 from uber_compose.uber_compose import TheUberCompose
 from uber_compose.vedro_plugin.base_structures.common_json_cli import CommandResult
@@ -27,5 +44,9 @@ __all__ = (
     'CommonJsonCli', 'CommandResult', 'ProcessExit', 'JsonParser', 'json_parser', 'ExecTimeout', 'TimeOutCheck',
     'ExecLifeCyclePolicy',
     'VedroUberCompose', 'UpHealthPolicy', 'DEFAULT_COMPOSE', 'ComposeConfig', 'DEFAULT_ENV_DESCRIPTION',
-    'DEFAULT_ENV_ID',
+    'DEFAULT_ENV_ID', 'UberComposeEvent',
+    'ExecutionPlan', 'WorkerAssignment', 'env_group_key', 'env_group_key_for',
+    'get_execution_plan', 'get_namespace_map', 'get_worker_groups', 'get_worker_id', 'get_workers_total',
+    'is_parallel_worker', 'is_parallel_worker_process', 'reset_execution_plan', 'set_execution_plan',
+    'VedroUberComposeParallel', 'build_execution_plan', 'build_execution_plan_from_scheduler',
 )

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/core/constants.py

@@ -6,11 +6,14 @@ class Constants:
     def __init__(self,
                  docker_compose_files_scan_depth: int = 2,
                  cli_compose_util_override: str = None,
+                 compose_project_name: str = None,
+                 tmp_envs_path: str | Path = None,
+                 in_docker_project_root_path: str | Path = None,
                  ):
         self.default_log_policy = os.environ.get('LOG_POLICY', 'DEFAULT')
 
-        self.project: str = os.environ.get('COMPOSE_PROJECT_NAME')
-        self.compose_project_name = os.environ.get('COMPOSE_PROJECT_NAME')
+        self.project: str = compose_project_name or os.environ.get('COMPOSE_PROJECT_NAME')
+        self.compose_project_name = self.project
         assert self.compose_project_name, (
             'COMPOSE_PROJECT_NAME environment variable is not set: - COMPOSE_PROJECT_NAME=${PWD##*/}'
         )
@@ -19,8 +22,12 @@ class Constants:
 
         # inner directories
         self.non_stop_containers: list[str] = os.environ.get('NON_STOP_CONTAINERS', 'e2e,dockersock').split(',')
-        self.tmp_envs_path: Path = Path(os.environ.get('TMP_ENVS_DIRECTORY', '/tmp/uc-envs'))
-        self.in_docker_project_root_path: Path = Path(os.environ.get('PROJECT_ROOT_DIRECTORY', '/project'))
+        self.tmp_envs_path: Path = Path(
+            tmp_envs_path or os.environ.get('TMP_ENVS_DIRECTORY', '/tmp/uc-envs')
+        )
+        self.in_docker_project_root_path: Path = Path(
+            in_docker_project_root_path or os.environ.get('PROJECT_ROOT_DIRECTORY', '/project')
+        )
 
         self.host_project_root_directory: Path = Path(
             os.environ.get('HOST_PROJECT_ROOT_DIRECTORY', '__host_project_root__')

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/core/docker_compose.py

@@ -17,6 +17,7 @@ from uber_compose.env_description.env_types import EventStage
 from uber_compose.errors.up import ServicesUpError
 from uber_compose.helpers.health_policy import UpHealthPolicy
 from uber_compose.helpers.jobs_result import JobResult
+from uber_compose.helpers.labels import Label
 from uber_compose.output.console import Logger
 from uber_compose.output.styles import Style
 
@@ -206,6 +207,7 @@ class ComposeInstance:
                 delay_s=self.health_policy.service_up_check_delay_s,
                 logger_func=self.logger.stage_details,
                 get_compose_state=self.compose_executor.dc_state,
+                service_filter=self._is_service_in_this_namespace,
             )
             if check_up_result != JobResult.GOOD:
                 services_status = await self.compose_executor.dc_state()
@@ -222,6 +224,12 @@ class ComposeInstance:
             migrations
         )
 
+    def _is_service_in_this_namespace(self, service_state) -> bool:
+        return (
+            service_state.labels.get(Label.ENV_ID)
+            == self.compose_instance_files.env_config_instance.env_id
+        )
+
     async def run(self):
         self.compose_instance_files = await self.generate_config_files()
 
@@ -242,6 +250,7 @@ class ComposeInstance:
             service
             for service_tier_pack in services_tiers
             for service in service_tier_pack
+            if service not in self.except_containers
         ]
 
         await self.run_migration(
@@ -267,6 +276,7 @@ class ComposeInstance:
                 delay_s=self.health_policy.service_up_check_delay_s,
                 logger_func=self.logger.stage,
                 get_compose_state=self.compose_executor.dc_state,
+                service_filter=self._is_service_in_this_namespace,
             )
             if check_up_result != JobResult.GOOD:
                 services_status = await self.compose_executor.dc_state()

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/core/sequence_run_types.py

@@ -12,6 +12,8 @@ class ComposeConfig(NamedTuple):
     compose_files: str
     overridden_services: list[OverridenService] = None
     parallel_env_limit: Union[int] = 1
+    env_parallelism: dict[str, int] | None = None
+    default_env_parallelism: int = 1
 
 
 class EnvInstanceConfig(NamedTuple):

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/core/system_docker_compose.py

@@ -107,5 +107,8 @@ class SystemDockerCompose:
         services_state = await self.dc_shell.dc_state()
         return [service.name for service in services_state.get_all_for()]
 
+    async def get_all_services_state(self) -> ServicesComposeState:
+        return await self.dc_shell.dc_state()
+
     async def down_services(self, services: list[str] | None = None) -> None:
         await self.dc_shell.dc_down(services=services)

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/core/utils/compose_instance_cfg.py

@@ -7,6 +7,8 @@ from uber_compose.env_description.env_types import Env
 from uber_compose.env_description.env_types import Environment
 from uber_compose.env_description.env_types import Service
 from uber_compose.env_description.env_types import ServiceMode
+from uber_compose.helpers.env_id import get_stable_env_id as get_stable_env_id  # noqa: F401
+from uber_compose.helpers.env_id import get_template_hash as get_template_hash  # noqa: F401
 
 
 def make_env_service_name(service, env_id):
@@ -52,8 +54,13 @@ def get_new_env_id() -> str:
 
 def get_service_map(env: Environment, new_env_id: str):
     return {
-        service_name: make_env_service_name(service_name, new_env_id)
-        for service_name, service in env.get_services().items() if service.mode != ServiceMode.OFF
+        service_name: (
+            service_name
+            if service.mode in (ServiceMode.OFF, ServiceMode.SINGLETON)
+            else make_env_service_name(service_name, new_env_id)
+        )
+        for service_name, service in env.get_services().items()
+        if service.mode != ServiceMode.OFF
     }
 
 

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/core/utils/state_waiting.py

@@ -23,10 +23,14 @@ def is_service_not_running_or_not_healthy(service_state: ServiceComposeState) ->
 
 
 async def wait_all_services_up(attempts, delay_s, logger_func: Callable, get_compose_state,
-                               is_service_in_bad_state=is_service_not_running_or_not_healthy) -> JobResult:
+                               is_service_in_bad_state=is_service_not_running_or_not_healthy,
+                               service_filter: Callable = None) -> JobResult:
     counter_keeper = CountdownCounterKeeper(attempts)
     state_keeper = StateKeeper()
 
+    if service_filter is None:
+        service_filter = lambda service_state: True
+
     async def check_once() -> JobResult:
         output_style = Style()
 
@@ -34,7 +38,7 @@ async def wait_all_services_up(attempts, delay_s, logger_func: Callable, get_com
             logger_func(Text('Starting services check up', style=output_style.info))
             state_keeper.update_state(ServicesState.DEFAULT_STATE)
 
-        services_state = await get_compose_state()
+        services_state = (await get_compose_state()).get_all_for(service_filter)
         non_ready_services = [
             service
             for service in services_state

{uber_compose-2.3.3 → uber_compose-3.0.0a1}/uber_compose/errors/up.py

@@ -5,3 +5,7 @@ class ServicesUpError(Exception):
 
     def __str__(self):
         return self.message
+
+
+class SingletonConfigChangedError(ServicesUpError):
+    ...

uber_compose-3.0.0a1/uber_compose/helpers/broken_services.py

@@ -0,0 +1,47 @@
+from uber_compose import Environment
+from uber_compose.core.docker_compose_shell.types import ServicesComposeState
+from uber_compose.core.utils.compose_instance_cfg import get_service_map
+from uber_compose.core.utils.state_waiting import is_service_not_running_or_not_healthy
+from uber_compose.helpers.labels import Label
+
+
+def calc_broken_services(
+    services_state: ServicesComposeState,
+    config_template: Environment,
+    excluded_from_check: list[str] = None,
+    env_id: str = None,
+) -> list[str]:
+    if excluded_from_check is None:
+        excluded_from_check = []
+
+    if env_id is None and services_state.get_any() is not None:
+        env_id = services_state.get_any().labels.get(Label.ENV_ID)
+
+    if env_id is not None:
+        expected_services_map = get_service_map(config_template, env_id)
+    else:
+        expected_services_map = {
+            service_name: service_name
+            for service_name in config_template.get_services_names()
+        }
+
+    non_ready_services = [
+        service.name
+        for service in services_state
+        if is_service_not_running_or_not_healthy(service)
+    ]
+    not_included_services = [
+        service
+        for service in expected_services_map.values()
+        if service not in services_state.get_services_names()
+    ]
+    skipped_overridden_services = [
+        expected_services_map.get(service.service.name, service.service.name)
+        for service in config_template.get_overridden_services()
+        if expected_services_map.get(service.service.name, service.service.name)
+        not in services_state.get_services_names()
+    ]
+    return list(
+        set(non_ready_services) | set(not_included_services)
+        - set(excluded_from_check) - set(skipped_overridden_services)
+    )

uber_compose-3.0.0a1/uber_compose/helpers/current_env.py

@@ -0,0 +1,20 @@
+from contextvars import ContextVar
+
+from uber_compose.core.sequence_run_types import DEFAULT_ENV_ID
+
+_current_env_id: ContextVar[str] = ContextVar(
+    'uber_compose_current_env_id',
+    default=DEFAULT_ENV_ID,
+)
+
+
+def set_current_env_id(env_id: str) -> None:
+    _current_env_id.set(env_id)
+
+
+def get_current_env_id() -> str:
+    return _current_env_id.get()
+
+
+def reset_current_env_id() -> None:
+    _current_env_id.set(DEFAULT_ENV_ID)

uber_compose-3.0.0a1/uber_compose/helpers/env_id.py

@@ -0,0 +1,21 @@
+import hashlib
+
+from uber_compose.helpers.bytes_pickle import base64_pickled
+
+
+def get_template_hash(config_template, length: int = 4) -> str:
+    payload = base64_pickled(config_template)
+    return hashlib.sha1(payload.encode()).hexdigest()[:length]
+
+
+def get_stable_env_id(config_template, index: int) -> str:
+    return f'w{index}-{get_template_hash(config_template)}'
+
+
+def stable_index(key: str, total: int) -> int:
+    total = max(total, 1)
+    digest = hashlib.sha1(str(key).encode()).hexdigest()
+    return int(digest, 16) % total
+
+
+__all__ = ("get_stable_env_id", "get_template_hash", "stable_index",)

uber_compose-3.0.0a1/uber_compose/helpers/hooks.py

@@ -0,0 +1,24 @@
+from collections import defaultdict
+from enum import Enum
+import inspect
+from typing import Any
+from typing import Callable
+
+
+class UberComposeEvent(Enum):
+    AFTER_UP = 'after_up'
+    AFTER_UPDATED_ENV_UP = 'after_updated_env_up'
+
+
+class HooksRegistry:
+    def __init__(self) -> None:
+        self._hooks: dict[UberComposeEvent, list[tuple[Callable[..., Any], tuple, dict]]] = defaultdict(list)
+
+    def on(self, event: UberComposeEvent, func: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
+        self._hooks[event].append((func, args, kwargs))
+
+    async def run(self, event: UberComposeEvent) -> None:
+        for func, args, kwargs in self._hooks[event]:
+            result = func(*args, **kwargs)
+            if inspect.isawaitable(result):
+                await result

uber_compose-3.0.0a1/uber_compose/parallel_plugin/__init__.py

@@ -0,0 +1,39 @@
+from uber_compose.parallel_plugin.grouper import assign_groups_to_workers
+from uber_compose.parallel_plugin.grouper import build_execution_plan
+from uber_compose.parallel_plugin.grouper import build_execution_plan_from_scheduler
+from uber_compose.parallel_plugin.grouper import build_replica_execution_plan
+from uber_compose.parallel_plugin.grouper import collect_group_counts
+from uber_compose.parallel_plugin.grouping import env_group_key
+from uber_compose.parallel_plugin.grouping import env_group_key_for
+from uber_compose.parallel_plugin.grouping import scenario_env
+from uber_compose.parallel_plugin.plugin import VedroUberComposeParallel
+from uber_compose.parallel_plugin.plugin import VedroUberComposeParallelPlugin
+from uber_compose.parallel_plugin.types import ExecutionPlan
+from uber_compose.parallel_plugin.types import WorkerAssignment
+from uber_compose.parallel_plugin.worker_api import get_env_parallelism
+from uber_compose.parallel_plugin.worker_api import get_env_replicas_total
+from uber_compose.parallel_plugin.worker_api import get_execution_plan
+from uber_compose.parallel_plugin.worker_api import get_namespace_map
+from uber_compose.parallel_plugin.worker_api import get_worker_groups
+from uber_compose.parallel_plugin.worker_api import get_worker_id
+from uber_compose.parallel_plugin.worker_api import get_worker_replica_index
+from uber_compose.parallel_plugin.worker_api import get_workers_total
+from uber_compose.parallel_plugin.worker_api import is_parallel_worker
+from uber_compose.parallel_plugin.worker_api import is_parallel_worker_process
+from uber_compose.parallel_plugin.worker_api import reset_env_parallelism
+from uber_compose.parallel_plugin.worker_api import reset_execution_plan
+from uber_compose.parallel_plugin.worker_api import set_env_parallelism
+from uber_compose.parallel_plugin.worker_api import set_execution_plan
+from uber_compose.parallel_plugin.worker_filter import ignore_foreign_scenarios
+
+__all__ = (
+    'ExecutionPlan', 'WorkerAssignment',
+    'VedroUberComposeParallel', 'VedroUberComposeParallelPlugin',
+    'assign_groups_to_workers', 'build_execution_plan', 'build_execution_plan_from_scheduler',
+    'build_replica_execution_plan', 'collect_group_counts', 'env_group_key', 'env_group_key_for',
+    'get_env_parallelism', 'get_env_replicas_total', 'get_execution_plan', 'get_namespace_map',
+    'get_worker_groups', 'get_worker_id', 'get_worker_replica_index', 'get_workers_total',
+    'ignore_foreign_scenarios', 'is_parallel_worker', 'is_parallel_worker_process',
+    'reset_env_parallelism', 'reset_execution_plan', 'scenario_env',
+    'set_env_parallelism', 'set_execution_plan',
+)

uber_compose-3.0.0a1/uber_compose/parallel_plugin/grouper.py

@@ -0,0 +1,97 @@
+from collections import defaultdict
+from typing import Callable
+
+from vedro.core import ScenarioScheduler
+from vedro.core import VirtualScenario
+
+from uber_compose.helpers.env_id import get_stable_env_id
+from uber_compose.parallel_plugin.types import ExecutionPlan
+from uber_compose.parallel_plugin.types import WorkerAssignment
+
+
+async def collect_group_counts(
+    scenarios: ScenarioScheduler,
+    grouping_key: Callable[[VirtualScenario], str],
+) -> dict[str, int]:
+    counts: dict[str, int] = defaultdict(int)
+    async for scenario in scenarios:
+        if scenario.is_skipped():
+            continue
+        counts[grouping_key(scenario)] += 1
+    return dict(counts)
+
+
+def assign_groups_to_workers(group_counts: dict[str, int], workers_total: int) -> list[list[str]]:
+    workers_total = max(workers_total, 1)
+    workers: list[list[str]] = [[] for _ in range(workers_total)]
+    worker_loads = [0] * workers_total
+
+    groups = sorted(group_counts.items(), key=lambda item: (-item[1], item[0]))
+    for group_key, count in groups:
+        worker_id = min(range(workers_total), key=lambda index: (worker_loads[index], index))
+        workers[worker_id].append(group_key)
+        worker_loads[worker_id] += count
+
+    return workers
+
+
+def build_execution_plan(
+    group_counts: dict[str, int],
+    workers_total: int,
+    namespace_map: dict[str, str] | None = None,
+) -> ExecutionPlan:
+    namespace_map = namespace_map or {}
+    workers = assign_groups_to_workers(group_counts, workers_total)
+    assignments = tuple(
+        WorkerAssignment(
+            worker_id=worker_id,
+            groups=tuple(sorted(groups)),
+            namespace_map={group: namespace_map[group] for group in groups if group in namespace_map},
+        )
+        for worker_id, groups in enumerate(workers)
+    )
+    return ExecutionPlan(workers_total=max(workers_total, 1), assignments=assignments)
+
+
+async def build_execution_plan_from_scheduler(
+    scenarios: ScenarioScheduler,
+    workers_total: int,
+    grouping_key: Callable[[VirtualScenario], str],
+    namespace_map: dict[str, str] | None = None,
+) -> ExecutionPlan:
+    group_counts = await collect_group_counts(scenarios, grouping_key)
+    return build_execution_plan(group_counts, workers_total, namespace_map)
+
+
+def build_replica_execution_plan(
+    group_counts: dict[str, int],
+    group_envs: dict[str, object],
+    namespace_counts: dict[str, int],
+    default_replicas: int = 1,
+) -> ExecutionPlan:
+    assignments = []
+    worker_id = 0
+    for group_key in sorted(group_envs):
+        env = group_envs[group_key]
+        replicas = max(int(namespace_counts.get(group_key, default_replicas)), 1)
+        replicas = min(replicas, max(group_counts.get(group_key, 1), 1))
+        for replica_index in range(replicas):
+            assignments.append(WorkerAssignment(
+                worker_id=worker_id,
+                groups=(group_key,),
+                namespace_map={group_key: get_stable_env_id(env, replica_index)},
+                replica_index=replica_index,
+                replicas_total=replicas,
+            ))
+            worker_id += 1
+    return ExecutionPlan(
+        workers_total=worker_id,
+        assignments=tuple(assignments),
+        env_templates=dict(group_envs),
+    )
+
+
+__all__ = (
+    'assign_groups_to_workers', 'build_execution_plan', 'build_execution_plan_from_scheduler',
+    'build_replica_execution_plan', 'collect_group_counts',
+)

uber_compose-3.0.0a1/uber_compose/parallel_plugin/grouping.py

@@ -0,0 +1,30 @@
+from vedro.core import VirtualScenario
+
+from uber_compose.env_description.env_types import DEFAULT_ENV_DESCRIPTION
+from uber_compose.env_description.env_types import Environment
+from uber_compose.helpers.bytes_pickle import base64_pickled
+
+
+def scenario_env(scenario: VirtualScenario) -> Environment | None:
+    original = getattr(scenario, '_orig_scenario', scenario)
+    env = None
+    if hasattr(original, 'tags'):
+        for tag in original.tags:
+            if isinstance(tag, Environment):
+                env = tag
+    if hasattr(original, 'env'):
+        env = original.env
+    return env
+
+
+def env_group_key_for(env: Environment | None) -> str:
+    if env is None or env.description == DEFAULT_ENV_DESCRIPTION:
+        env = None
+    return base64_pickled(env)
+
+
+def env_group_key(scenario: VirtualScenario) -> str:
+    return env_group_key_for(scenario_env(scenario))
+
+
+__all__ = ("env_group_key", "env_group_key_for", "scenario_env",)

uber_compose-3.0.0a1/uber_compose/parallel_plugin/plugin.py

@@ -0,0 +1,149 @@
+from typing import Type
+
+from vedro.core import ConfigType
+from vedro.core import Dispatcher
+from vedro.core import Plugin
+from vedro.core import PluginConfig
+from vedro.events import ArgParseEvent
+from vedro.events import ArgParsedEvent
+from vedro.events import ConfigLoadedEvent
+from vedro.events import StartupEvent
+
+from uber_compose.env_description.env_types import DEFAULT_ENV_DESCRIPTION
+from uber_compose.parallel_plugin.grouper import build_execution_plan_from_scheduler
+from uber_compose.parallel_plugin.grouper import build_replica_execution_plan
+from uber_compose.parallel_plugin.grouper import collect_group_counts
+from uber_compose.parallel_plugin.grouping import env_group_key
+from uber_compose.parallel_plugin.grouping import env_group_key_for
+from uber_compose.parallel_plugin.grouping import scenario_env
+from uber_compose.parallel_plugin.runner import ParallelScenarioRunner
+from uber_compose.parallel_plugin.worker_api import get_env_parallelism
+from uber_compose.parallel_plugin.worker_api import get_env_replicas_total
+from uber_compose.parallel_plugin.worker_api import get_worker_groups
+from uber_compose.parallel_plugin.worker_api import get_worker_replica_index
+from uber_compose.parallel_plugin.worker_api import is_parallel_worker_process
+from uber_compose.parallel_plugin.worker_api import set_execution_plan
+from uber_compose.parallel_plugin.worker_filter import ignore_foreign_scenarios
+
+
+class VedroUberComposeParallelPlugin(Plugin):
+    def __init__(self, config: Type["VedroUberComposeParallel"]) -> None:
+        super().__init__(config)
+        self._enabled = config.enabled
+        self._grouping_key = config.grouping_key or env_group_key
+        self._parallel = max(int(config.parallel or 1), 1)
+        self._global_config: ConfigType | None = None
+        self._runner_registered = False
+        self._refresh_worker_state()
+
+    def _refresh_worker_state(self) -> None:
+        self._is_worker = is_parallel_worker_process()
+        self._worker_groups = get_worker_groups()
+        self._worker_replica_index = get_worker_replica_index()
+        self._env_replicas_total = get_env_replicas_total()
+
+    def subscribe(self, dispatcher: Dispatcher) -> None:
+        if not self._enabled:
+            return
+
+        dispatcher.listen(ConfigLoadedEvent, self.on_config_loaded) \
+            .listen(ArgParseEvent, self.handle_arg_parse, priority=-100) \
+            .listen(ArgParsedEvent, self.handle_arg_parsed, priority=-100) \
+            .listen(StartupEvent, self.handle_prepare_scenarios, priority=-100)
+
+    def on_config_loaded(self, event: ConfigLoadedEvent) -> None:
+        self._global_config = event.config
+
+    def handle_arg_parse(self, event: ArgParseEvent) -> None:
+        group = event.arg_parser.add_argument_group("Uber Compose Parallel")
+        group.add_argument("--uc-parallel",
+                           type=int,
+                           default=None,
+                           help="Number of parallel workers (1 keeps the sequential run)")
+
+    def handle_arg_parsed(self, event: ArgParsedEvent) -> None:
+        parallel = getattr(event.args, 'uc_parallel', None)
+        if parallel:
+            self._parallel = max(int(parallel), 1)
+        self._refresh_worker_state()
+
+    async def handle_prepare_scenarios(self, event: StartupEvent) -> None:
+        self._refresh_worker_state()
+
+        if self._is_worker:
+            await ignore_foreign_scenarios(
+                event.scheduler,
+                self._worker_groups,
+                self._grouping_key,
+                replica_index=self._worker_replica_index,
+                replicas_total=self._env_replicas_total,
+            )
+            return
+
+        plan = await self._build_plan(event)
+        if plan is None or plan.workers_total <= 1:
+            return
+
+        set_execution_plan(plan)
+        self._register_runner()
+
+    async def _build_plan(self, event: StartupEvent):
+        env_parallelism = get_env_parallelism()
+        if env_parallelism is not None:
+            return await self._build_replica_plan(event, env_parallelism)
+        if self._parallel <= 1:
+            return None
+        return await build_execution_plan_from_scheduler(
+            event.scheduler,
+            self._parallel,
+            self._grouping_key,
+        )
+
+    async def _build_replica_plan(self, event: StartupEvent, env_parallelism):
+        counts_by_description, default_replicas = env_parallelism
+        group_counts = await collect_group_counts(event.scheduler, self._grouping_key)
+        group_envs = {}
+        namespace_counts = {}
+
+        async for scenario in event.scheduler:
+            if scenario.is_skipped():
+                continue
+            env = scenario_env(scenario)
+            group_key = env_group_key_for(env)
+            template = None if env is None or env.description == DEFAULT_ENV_DESCRIPTION else env
+            group_envs[group_key] = template
+            description = template.description if template is not None else DEFAULT_ENV_DESCRIPTION
+            namespace_counts[group_key] = counts_by_description.get(description, default_replicas)
+
+        return build_replica_execution_plan(
+            group_counts=group_counts,
+            group_envs=group_envs,
+            namespace_counts=namespace_counts,
+            default_replicas=default_replicas,
+        )
+
+    def _register_runner(self) -> None:
+        if self._runner_registered or self._is_worker or self._global_config is None:
+            return
+        dispatcher = self._global_config.Registry.Dispatcher()
+        self._global_config.Registry.ScenarioRunner.register(
+            resolver=lambda: ParallelScenarioRunner(
+                dispatcher=dispatcher,
+                grouping_key=self._grouping_key,
+            ),
+            registrant=self,
+        )
+        self._runner_registered = True
+
+
+class VedroUberComposeParallel(PluginConfig):
+    plugin = VedroUberComposeParallelPlugin
+
+    enabled = False
+
+    parallel: int = 1
+
+    grouping_key = env_group_key
+
+
+__all__ = ("VedroUberComposeParallel", "VedroUberComposeParallelPlugin",)