From c0525703fb850d566edcb77deeae2940ab857f5a Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Mon, 23 Feb 2026 10:44:10 +0100 Subject: [PATCH 01/12] Adding hooks to parse out a resource manager requirements, updating a config to allow for testing --- config/tests/one-controller-config.data.xml | 6 +++++ src/drunc/unified_shell/shell.py | 25 +++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/config/tests/one-controller-config.data.xml b/config/tests/one-controller-config.data.xml index 7b3ee77c0..b5f745eab 100644 --- a/config/tests/one-controller-config.data.xml +++ b/config/tests/one-controller-config.data.xml @@ -105,6 +105,11 @@ + + + + + @@ -124,6 +129,7 @@ + diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index e6776a6af..1f6b37068 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -11,6 +11,7 @@ import click import click_shell import conffwk +import confmodel_dal from daqpytools.logging.levels import logging_log_levels from druncschema.description_pb2 import Description from druncschema.process_manager_pb2 import ProcessQuery @@ -177,6 +178,30 @@ def unified_shell( session_dal = db.get_dal(class_name="Session", uid=ctx.obj.configuration_id) app_log_path = session_dal.log_path + # Get access to the resource manager if it is defined in the configuration + resource_manager: "confmodel_dal.ResourceManagerConf | None" = None + if getattr(session_dal, "resource_manager", None): + resource_manager = session_dal.resource_manager + + # If the resource manager is present, extract its parameters and log them + resource_manager_address: str | None = None + if resource_manager: + resource_manager_address = f"{resource_manager.address}:{resource_manager.port}" + unified_shell_log.info( + f"Found resource manager at address {resource_manager_address}" + ) + + # If the resource manager is present, get the list of resources required from it + if resource_manager: + session_requested_resources: set(str) = ( + confmodel_dal.segment_get_managed_object_tags( + db._obj, ctx.obj.configuration_id + ) + ) + unified_shell_log.info( + f"This session has requested the following resources: {session_requested_resources=}" + ) + unified_shell_log.info( f"[green]Setting up to use the process manager[/green] with configuration " f"[green]{process_manager}[/green] and configuration id [green]" From 2ec5b0ec2e31240e7e0472a79d750e13b8833111 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Mon, 23 Feb 2026 11:20:37 +0100 Subject: [PATCH 02/12] Adding log record in the case that there is no session manager. --- src/drunc/unified_shell/shell.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index 1f6b37068..d5de2a659 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -182,6 +182,11 @@ def unified_shell( resource_manager: "confmodel_dal.ResourceManagerConf | None" = None if getattr(session_dal, "resource_manager", None): resource_manager = session_dal.resource_manager + else: + # This will be changed to an error once the resource manager service is deployed and all configurations have a resource manager configuration. + unified_shell_log.info( + f"The session {ctx.obj.configuration_id} in file {ctx.obj.configuration_file} does not have a resource manager defined in the file, resources will not be parsed out." + ) # If the resource manager is present, extract its parameters and log them resource_manager_address: str | None = None From f0c2bad3c57461c9da78cc93ed8c52ddd228fcb2 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Mon, 23 Feb 2026 11:23:28 +0100 Subject: [PATCH 03/12] Making the implication of a lack of a resource manager clearer --- src/drunc/unified_shell/shell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index d5de2a659..db70da8d8 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -185,7 +185,7 @@ def unified_shell( else: # This will be changed to an error once the resource manager service is deployed and all configurations have a resource manager configuration. unified_shell_log.info( - f"The session {ctx.obj.configuration_id} in file {ctx.obj.configuration_file} does not have a resource manager defined in the file, resources will not be parsed out." + f"The session [green]{ctx.obj.configuration_id} in file {ctx.obj.configuration_file}[/green] [yellow]does not have a resource manager defined[/yellow], resources will not be requested from the resource management service" ) # If the resource manager is present, extract its parameters and log them From 214297cbaaff86a3bb20bc00cd1535715b407d35 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Mon, 23 Feb 2026 15:39:02 +0100 Subject: [PATCH 04/12] Removing the resource manager configuration - the requirements are not yet understood --- src/drunc/unified_shell/shell.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index db70da8d8..185683308 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -188,14 +188,6 @@ def unified_shell( f"The session [green]{ctx.obj.configuration_id} in file {ctx.obj.configuration_file}[/green] [yellow]does not have a resource manager defined[/yellow], resources will not be requested from the resource management service" ) - # If the resource manager is present, extract its parameters and log them - resource_manager_address: str | None = None - if resource_manager: - resource_manager_address = f"{resource_manager.address}:{resource_manager.port}" - unified_shell_log.info( - f"Found resource manager at address {resource_manager_address}" - ) - # If the resource manager is present, get the list of resources required from it if resource_manager: session_requested_resources: set(str) = ( @@ -204,7 +196,7 @@ def unified_shell( ) ) unified_shell_log.info( - f"This session has requested the following resources: {session_requested_resources=}" + f"This session has requested the following resources: {session_requested_resources}" ) unified_shell_log.info( From c63b360e0333c7e098d2b77a4b69d4154685bc74 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Mon, 23 Feb 2026 17:14:00 +0100 Subject: [PATCH 05/12] Parsing out the per-segment requirements from the config --- src/drunc/unified_shell/shell.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index 185683308..564ba0870 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -188,6 +188,22 @@ def unified_shell( f"The session [green]{ctx.obj.configuration_id} in file {ctx.obj.configuration_file}[/green] [yellow]does not have a resource manager defined[/yellow], resources will not be requested from the resource management service" ) + # Iterate through all the segment nest levels, parse out the requested managed + # objects for that segment, and allocate them to a dict + managed_objects: dict[ + str : list(str) + ] = {} # segment: list[managed_object_identifier] + segments = session_dal.segment.segments + while segments: + nested_segments = [] + for segment in segments: + managed_objects[segment.id] = confmodel_dal.segment_get_managed_object_tags( + db._obj, ctx.obj.configuration_id, segment.id + ) + nested_segments += [nested_segment for nested_segment in segment.segments] + segments = nested_segments + unified_shell_log.info(f"{managed_objects=}") + # If the resource manager is present, get the list of resources required from it if resource_manager: session_requested_resources: set(str) = ( From c10a7a67da18d61c92b84feef72dd785768afe82 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Thu, 26 Feb 2026 20:01:06 +0100 Subject: [PATCH 06/12] Working request allocation logging --- config/tests/deep-segments-config.data.xml | 1 + config/tests/nestedConfig.data.xml | 1 + src/drunc/unified_shell/commands.py | 84 ++++++++++++++++++++++ src/drunc/unified_shell/context.py | 4 ++ src/drunc/unified_shell/shell.py | 62 +++++----------- src/drunc/unified_shell/shell_utils.py | 18 +++++ 6 files changed, 126 insertions(+), 44 deletions(-) diff --git a/config/tests/deep-segments-config.data.xml b/config/tests/deep-segments-config.data.xml index 796e32b0f..d2c40434e 100644 --- a/config/tests/deep-segments-config.data.xml +++ b/config/tests/deep-segments-config.data.xml @@ -226,6 +226,7 @@ + diff --git a/config/tests/nestedConfig.data.xml b/config/tests/nestedConfig.data.xml index 23592ceb8..2102672cd 100644 --- a/config/tests/nestedConfig.data.xml +++ b/config/tests/nestedConfig.data.xml @@ -217,6 +217,7 @@ + diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index 4bfc76cfe..6b527a287 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -2,11 +2,14 @@ import sys import click +import conffwk +import confmodel_dal from druncschema.process_manager_pb2 import ProcessQuery from drunc.controller.interface.shell_utils import controller_setup from drunc.process_manager.interface.context import ProcessManagerContext from drunc.unified_shell.context import UnifiedShellMode +from drunc.unified_shell.shell_utils import resource_log_tree from drunc.utils.shell_utils import InterruptedCommand from drunc.utils.utils import get_logger @@ -25,14 +28,58 @@ help="Sleep between app boot, in seconds. This may be useful if you have are using SSHPM, and have SSHD's maxstartups setting set to a low value.", ) @click.pass_obj +@click.pass_context def boot( + ctx: click.core.Context, obj: ProcessManagerContext, override_logs: bool | None, sleep_between_app_boot: int | float = 0, ) -> None: log = get_logger("unified_shell.boot") + + # Instantiate the session dal to parse out the managed objects + db = conffwk.Configuration(ctx.obj.configuration_file) + session_dal = db.get_dal(class_name="Session", uid=ctx.obj.configuration_id) session_name = obj.session_name user = getpass.getuser() + + # Iterate through all the segment nest levels, parse out the requested managed + # objects for that segment, and allocate them to a dict + managed_objects: dict[ + str : list(str) + ] = {} # segment: list[managed_object_identifier] + managed_objects_present: bool = False + segments = session_dal.segment.segments + while segments: + nested_segments = [] + for segment in segments: + managed_objects[segment.id] = confmodel_dal.segment_get_managed_object_tags( + db._obj, ctx.obj.configuration_id, segment.id + ) + if managed_objects[segment.id]: + managed_objects_present = True + nested_segments += [nested_segment for nested_segment in segment.segments] + segments = nested_segments + ctx.obj.managed_objects_present = managed_objects_present + ctx.obj.managed_objects = managed_objects + + # Split out the segments that have requested resources + empty_segments = [k for k, v in managed_objects.items() if not v] + active_segments = {k: v for k, v in managed_objects.items() if v} + + # Log the request of resources if they are used + if ctx.obj.managed_objects_present: + log.info( + "[blue]Placeholder[/blue] Requesting objects in the following segments:" + ) + # Note the next 4 lines should be considered to be indented + if active_segments: + resource_log_tree(active_segments, log) + if empty_segments: + log.info( + f"[yellow]Empty segments (skipped):[/yellow] {', '.join(empty_segments)}" + ) + processes = obj.get_driver("process_manager").ps( ProcessQuery(user=user, session=session_name) ) @@ -97,6 +144,43 @@ def boot( sys.exit(1) +@click.command("terminate") +@click.pass_obj +@click.pass_context +def terminate(ctx, obj): + """ + Execute the process manager terminate command, but release the resources prior to + doing so + """ + + log = get_logger("unified_shell.terminate") + + # Get the handle to the managed objects + all_objects = ctx.obj.managed_objects + + # Split out the segments that have requested resources + empty_segments = [k for k, v in all_objects.items() if not v] + active_segments = {k: v for k, v in all_objects.items() if v} + + # Log the release of requested resources if they were used + if ctx.obj.managed_objects_present: + log.info( + "[blue]Placeholder[/blue] Releasing managed objects in the following segments:" + ) + + # if ctx.obj.managed_objects_present:all_objects = ctx.obj.managed_objects + if active_segments: + resource_log_tree(active_segments, log) + if empty_segments: + log.info( + f"[yellow]Empty segments (skipped):[/yellow] {', '.join(empty_segments)}" + ) + ctx.obj.managed_objects = {} + ctx.obj.managed_objects_present = False + + obj.get_driver("process_manager").terminate() + + @click.command("start-shell") @click.pass_obj @click.pass_context diff --git a/src/drunc/unified_shell/context.py b/src/drunc/unified_shell/context.py index 318983739..963185ab2 100644 --- a/src/drunc/unified_shell/context.py +++ b/src/drunc/unified_shell/context.py @@ -25,6 +25,10 @@ def __init__(self): self.session_name = "" self.override_logs = True self.running_mode = UnifiedShellMode.INTERACTIVE + self.managed_objects: dict[ + str : list(str) + ] = {} # segment: list[managed_object_identifier] + self.managed_objects_present: bool = False super(UnifiedShellContext, self).__init__() def reset(self, address_pm: str = ""): diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index 564ba0870..78baa6409 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -11,7 +11,6 @@ import click import click_shell import conffwk -import confmodel_dal from daqpytools.logging.levels import logging_log_levels from druncschema.description_pb2 import Description from druncschema.process_manager_pb2 import ProcessQuery @@ -49,11 +48,10 @@ logs, ps, restart, - terminate, ) from drunc.process_manager.interface.process_manager import run_pm from drunc.process_manager.utils import get_pm_type_from_name, validate_k8s_session_name -from drunc.unified_shell.commands import boot, start_shell +from drunc.unified_shell.commands import boot, start_shell, terminate from drunc.unified_shell.context import UnifiedShellMode from drunc.unified_shell.shell_utils import generate_fsm_sequence_command from drunc.utils.configuration import ConfTypes, OKSKey @@ -178,43 +176,6 @@ def unified_shell( session_dal = db.get_dal(class_name="Session", uid=ctx.obj.configuration_id) app_log_path = session_dal.log_path - # Get access to the resource manager if it is defined in the configuration - resource_manager: "confmodel_dal.ResourceManagerConf | None" = None - if getattr(session_dal, "resource_manager", None): - resource_manager = session_dal.resource_manager - else: - # This will be changed to an error once the resource manager service is deployed and all configurations have a resource manager configuration. - unified_shell_log.info( - f"The session [green]{ctx.obj.configuration_id} in file {ctx.obj.configuration_file}[/green] [yellow]does not have a resource manager defined[/yellow], resources will not be requested from the resource management service" - ) - - # Iterate through all the segment nest levels, parse out the requested managed - # objects for that segment, and allocate them to a dict - managed_objects: dict[ - str : list(str) - ] = {} # segment: list[managed_object_identifier] - segments = session_dal.segment.segments - while segments: - nested_segments = [] - for segment in segments: - managed_objects[segment.id] = confmodel_dal.segment_get_managed_object_tags( - db._obj, ctx.obj.configuration_id, segment.id - ) - nested_segments += [nested_segment for nested_segment in segment.segments] - segments = nested_segments - unified_shell_log.info(f"{managed_objects=}") - - # If the resource manager is present, get the list of resources required from it - if resource_manager: - session_requested_resources: set(str) = ( - confmodel_dal.segment_get_managed_object_tags( - db._obj, ctx.obj.configuration_id - ) - ) - unified_shell_log.info( - f"This session has requested the following resources: {session_requested_resources}" - ) - unified_shell_log.info( f"[green]Setting up to use the process manager[/green] with configuration " f"[green]{process_manager}[/green] and configuration id [green]" @@ -345,8 +306,10 @@ def unified_shell( # Add the unified shell Click commands to the CLI unified_shell_log.debug("Adding [green]unified_shell[/green] commands") - ctx.command.add_command(boot, "boot") - ctx.obj.dynamic_commands.add("boot") + unified_shell_commands: list[click.Command] = [boot, terminate] + for cmd in unified_shell_commands: + ctx.command.add_command(cmd, format_name_for_cli(cmd.name)) + ctx.obj.dynamic_commands.add(format_name_for_cli(cmd.name)) # Add the process manager Click commands to the CLI unified_shell_log.debug("Adding [green]process_manager[/green] commands") @@ -451,7 +414,14 @@ def cleanup(): # Attempt a stateful shutdown of the controller if possible, returning to # initial state before terminating - if ctx.obj.get_driver("controller", quiet_fail=True): + if ( + len( + ctx.obj.get_driver("process_manager") + .ps(ProcessQuery(user=getpass.getuser(), session=ctx.obj.session_name)) + .values + ) + > 0 + ) and ctx.obj.get_driver("controller", quiet_fail=True): try: if ctx.obj.get_driver("controller").status().status.in_error: unified_shell_log.warning( @@ -493,7 +463,11 @@ def cleanup(): # Terminate any residual processes if ctx.obj.get_driver("process_manager"): - ctx.obj.get_driver("process_manager").terminate() + terminate_cmd = ctx.command.get_command(ctx, "terminate") + if terminate_cmd: + ctx.invoke(terminate_cmd) + else: + unified_shell_log.error("Command 'terminate' not found.") # Check if any processes are still running if ( diff --git a/src/drunc/unified_shell/shell_utils.py b/src/drunc/unified_shell/shell_utils.py index 389f4c542..38956c349 100644 --- a/src/drunc/unified_shell/shell_utils.py +++ b/src/drunc/unified_shell/shell_utils.py @@ -161,3 +161,21 @@ def generate_fsm_sequence_command( )(cmd) return cmd, format_name_for_cli(sequence.id) + + +def resource_log_tree(data, log, prefix=""): + items = list(data.items()) + total = len(items) + + for i, (key, value) in enumerate(items): + is_last = i == total - 1 + connector = "└── " if is_last else "├── " + + # Check if the value is a nested segment (resources stored as dict) + if isinstance(value, dict) and value: + extension = " " if is_last else "│ " + resource_log_tree(value, prefix + extension) + + else: + display_val = str(value) + log.info(f"{prefix}{connector}{key}: {display_val}") From 9b2b9db6558ddd270a05027956c26a01bba057bd Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Thu, 30 Apr 2026 13:58:24 +0200 Subject: [PATCH 07/12] Formatting the requested resources --- src/drunc/unified_shell/commands.py | 37 +++++++++++++++++++++++++++-- src/drunc/unified_shell/shell.py | 12 ++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index e3414ae91..130e067ce 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -1,4 +1,6 @@ import getpass +import os +import socket import sys import click @@ -54,8 +56,10 @@ def boot( while segments: nested_segments = [] for segment in segments: - managed_objects[segment.id] = confmodel_dal.segment_get_managed_object_tags( - db._obj, ctx.obj.configuration_id, segment.id + managed_objects[segment.id] = list( + confmodel_dal.segment_get_managed_object_tags( + db._obj, ctx.obj.configuration_id, segment.id + ) ) if managed_objects[segment.id]: managed_objects_present = True @@ -64,6 +68,35 @@ def boot( ctx.obj.managed_objects_present = managed_objects_present ctx.obj.managed_objects = managed_objects + # Map the requested dataflow localhost paths to realpaths, and localhost to host names + for segment, _managed_objects in managed_objects.items(): + log.critical( + f"Segment '{segment}' has requested the following managed objects: {', '.join(_managed_objects)}" + ) + for i, managed_object in enumerate(_managed_objects): + # Correct the storage paths if necessary + if managed_object.startswith("storage:"): + log.debug(f"Mapping storage path '{managed_object}' to real path") + + # Map localhost to the host name + if "localhost" in managed_object: + updated_host = managed_object.replace( + "localhost", socket.gethostname() + ) + _managed_objects[i] = updated_host + + # Map the path to a real path, the paths are commonly "." + parts = _managed_objects[i].split(":") + raw_path = parts[-1] + real_path = os.path.abspath(raw_path) + mount = "/".join(real_path.split("/")[:2]) + + prefix = ":".join(parts[:-1]) + _managed_objects[i] = f"{prefix}:{mount}" + log.critical( + f"Mapped storage path '{managed_object}' to real path '{_managed_objects[i]}'" + ) + # Split out the segments that have requested resources empty_segments = [k for k, v in managed_objects.items() if not v] active_segments = {k: v for k, v in managed_objects.items() if v} diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index a3bc1122e..258c2665c 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -183,6 +183,18 @@ def unified_shell( session_dal = db.get_dal(class_name="Session", uid=ctx.obj.configuration_id) app_log_path = session_dal.log_path + # Get the session manager URL - FOR DEV + resource_manager_url = getattr(session_dal, "resource_manager", None) + if not resource_manager_url: + unified_shell_log.info("No resource manager URL found in the configuration.") + else: + resource_manager_host = resource_manager_url.address + resource_manager_port = resource_manager_url.port + resource_manager_url = f"http://{resource_manager_host}:{resource_manager_port}" + unified_shell_log.info( + f"Resource manager URL from configuration: [green]{resource_manager_url}[/green]" + ) + unified_shell_log.info( f"[green]Setting up to use the process manager[/green] with configuration " f"[green]{process_manager}[/green] and configuration id [green]" From 0e4c2f961f8ff6f63ecb4bed1b77b863df7e4951 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Thu, 30 Apr 2026 16:47:26 +0200 Subject: [PATCH 08/12] Working first demo --- request_test.py | 14 +++ src/drunc/resource_manager/__init__.py | 5 + src/drunc/resource_manager/client.py | 122 +++++++++++++++++++++++++ src/drunc/unified_shell/commands.py | 31 ++++++- src/drunc/unified_shell/context.py | 3 + src/drunc/unified_shell/shell.py | 4 +- 6 files changed, 177 insertions(+), 2 deletions(-) create mode 100644 request_test.py create mode 100644 src/drunc/resource_manager/__init__.py create mode 100644 src/drunc/resource_manager/client.py diff --git a/request_test.py b/request_test.py new file mode 100644 index 000000000..2fc75e378 --- /dev/null +++ b/request_test.py @@ -0,0 +1,14 @@ +import requests + +url = "http://127.0.0.1:8000/api/request_resource/" +payload = { + "names": "resource_one", + "owner": "pplesnia", + "session_id": "1234567890", + "session_name": "test_session", +} + +# verify=False mimics the -k flag in curl (disables SSL verification) +response = requests.post(url, data=payload, verify=False) +print(f"Status Code: {response.status_code}") +print(response.json()) diff --git a/src/drunc/resource_manager/__init__.py b/src/drunc/resource_manager/__init__.py new file mode 100644 index 000000000..0099763c6 --- /dev/null +++ b/src/drunc/resource_manager/__init__.py @@ -0,0 +1,5 @@ +from drunc.utils.utils import get_logger + +# Initialise process manager logger with Rich handler +# This is the tty interface, so its designed to be coloured +get_logger("resoruce_manager", rich_handler=True) diff --git a/src/drunc/resource_manager/client.py b/src/drunc/resource_manager/client.py new file mode 100644 index 000000000..c1838e1f1 --- /dev/null +++ b/src/drunc/resource_manager/client.py @@ -0,0 +1,122 @@ +import requests + +from drunc.utils.utils import get_logger + + +class ResourceManagerClient: + """ + Interface for communicating with the Resource Manager service. + """ + + def __init__(self, base_url): + """ + Initialize the ResourceManagerClient with the base URL of the Resource Manager service. + """ + self.url = base_url.rstrip("/") + self.log = get_logger("resource_manager.client") + + def _send_request(self, endpoint, payload): + """ + Helper method to send a POST request to the Resource Manager and handle responses. + + Args: + endpoint (str): The full URL endpoint to send the request to + payload (dict): The data payload to send in the request + + Returns: + dict: The JSON response from the server if successful, or None if an error occurred + + Raises: + None: Logs errors and returns None instead of raising exceptions for HTTP errors or unexpected issues. + """ + try: + # verify=False is used here for local/self-signed certs (like curl -k) + self.log.critical(f"Sending request to {endpoint} with payload: {payload}") + response = requests.post(endpoint, data=payload, verify=False) + + if "application/json" not in response.headers.get("Content-Type", ""): + self.log.error( + "Server returned HTML/Text instead of JSON. Check your URL paths." + ) + return None + + # Raise an exception for 4xx or 5xx status codes + response.raise_for_status() + + # Log the successful response + return response.json() + + except requests.exceptions.HTTPError: + self.log.warning(f"Request failed with status {response.status_code}") + return response.json() + except Exception as e: + self.log.error(f"An unexpected error occurred: {e}") + return None + + def query_resources( + self, resources: list[str], owner: str, session_id: str, session_name: str + ): + """ + Query the Resource Manager for the status of the specified resources. + + Args: + resources (list[str]): List of resource names to query + owner (str): The new owner of the resources + session_id (str): The session ID taking the resources + session_name (str): The session name taking the resources + + Returns: + dict: A dictionary containing the status of the queried resources + """ + payload = { + "names": ",".join(resources), + } + endpoint = f"{self.url}/api/query_resource/" + + return self._send_request(endpoint, payload) + + def request_resources( + self, resources: list[str], owner: str, session_id: str, session_name: str + ): + """ + Request resources from the Resource Manager for isolation during the run. + + Args: + resources (list[str]): List of resource names to query + owner (str): The new owner of the resources + session_id (str): The session ID taking the resources + session_name (str): The session name taking the resources + + Returns: + dict: A dictionary containing the status of the queried resources + """ + payload = { + "names": ",".join(resources), + "owner": owner, + "session_id": session_id, + "session_name": session_name, + } + endpoint = f"{self.url}/api/request_resource/" + + return self._send_request(endpoint, payload) + + def release_resources(self, resources: list[str], owner: str): + """ + Release resources from the Resource Manager for other runs to use. + + Args: + resources (list[str]): List of resource names to query + owner (str): The new owner of the resources + session_id (str): The session ID taking the resources + session_name (str): The session name taking the resources + + Returns: + dict: A dictionary containing the status of the queried resources + """ + payload = { + "names": ",".join(resources), + "owner": owner, + } + endpoint = f"{self.url}/api/release_resource/" + + return self._send_request(endpoint, payload) diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index 130e067ce..97788ad3b 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -52,15 +52,18 @@ def boot( str : list(str) ] = {} # segment: list[managed_object_identifier] managed_objects_present: bool = False + session_resources: list[str] = [] segments = session_dal.segment.segments while segments: nested_segments = [] for segment in segments: - managed_objects[segment.id] = list( + segment_resources = list( confmodel_dal.segment_get_managed_object_tags( db._obj, ctx.obj.configuration_id, segment.id ) ) + managed_objects[segment.id] = segment_resources + session_resources += segment_resources if managed_objects[segment.id]: managed_objects_present = True nested_segments += [nested_segment for nested_segment in segment.segments] @@ -114,6 +117,23 @@ def boot( f"[yellow]Empty segments (skipped):[/yellow] {', '.join(empty_segments)}" ) + # Remove storage related ones for initial prototyping + ctx.obj.session_resources = [ + r for r in session_resources if not r.startswith("storage:") + ] + + # Request the resources from the resource manager + if ctx.obj.resource_manager_client and ctx.obj.session_resources: + log.info( + f"Requesting the following resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {', '.join(ctx.obj.session_resources)}" + ) + ctx.obj.resource_manager_client.request_resources( + ctx.obj.session_resources, + getpass.getuser(), + ctx.obj.configuration_id, + session_name, + ) + processes = obj.get_driver("process_manager").ps( ProcessQuery(user=user, session=session_name) ) @@ -215,6 +235,15 @@ def terminate(ctx, obj): ctx.obj.managed_objects = {} ctx.obj.managed_objects_present = False + if ctx.obj.resource_manager_client and ctx.obj.session_resources: + log.info( + f"Releasing the following resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {', '.join(ctx.obj.session_resources)}" + ) + ctx.obj.resource_manager_client.release_resources( + ctx.obj.session_resources, getpass.getuser() + ) + ctx.obj.session_resources = [] + obj.get_driver("process_manager").terminate() diff --git a/src/drunc/unified_shell/context.py b/src/drunc/unified_shell/context.py index d3464721d..c984758f9 100644 --- a/src/drunc/unified_shell/context.py +++ b/src/drunc/unified_shell/context.py @@ -3,6 +3,7 @@ from druncschema.token_pb2 import Token +from drunc.resource_manager.client import ResourceManagerClient from drunc.utils.shell_utils import ShellContext @@ -30,6 +31,8 @@ def __init__(self): str : list(str) ] = {} # segment: list[managed_object_identifier] self.managed_objects_present: bool = False + self.resource_manager_client: ResourceManagerClient | None = None + self.session_resources: list[str] = [] super(UnifiedShellContext, self).__init__() def reset(self, address_pm: str = ""): diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index 258c2665c..0d1784097 100644 --- a/src/drunc/unified_shell/shell.py +++ b/src/drunc/unified_shell/shell.py @@ -58,6 +58,7 @@ ) from drunc.process_manager.interface.process_manager import run_pm from drunc.process_manager.utils import get_pm_type_from_name, validate_k8s_session_name +from drunc.resource_manager.client import ResourceManagerClient from drunc.unified_shell.commands import boot, start_shell, terminate from drunc.unified_shell.context import UnifiedShellMode from drunc.unified_shell.shell_utils import generate_fsm_sequence_command @@ -192,8 +193,9 @@ def unified_shell( resource_manager_port = resource_manager_url.port resource_manager_url = f"http://{resource_manager_host}:{resource_manager_port}" unified_shell_log.info( - f"Resource manager URL from configuration: [green]{resource_manager_url}[/green]" + f"Resource manager URL parsed from configuration: [green]{resource_manager_url}[/green]" ) + ctx.obj.resource_manager_client = ResourceManagerClient(resource_manager_url) unified_shell_log.info( f"[green]Setting up to use the process manager[/green] with configuration " From 493440bbbba933b735acb44b89ba115fc97d30d3 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Fri, 8 May 2026 17:32:14 +0200 Subject: [PATCH 09/12] Fixing some typos --- src/drunc/resource_manager/__init__.py | 2 +- src/drunc/unified_shell/commands.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/drunc/resource_manager/__init__.py b/src/drunc/resource_manager/__init__.py index 0099763c6..b56b7ee73 100644 --- a/src/drunc/resource_manager/__init__.py +++ b/src/drunc/resource_manager/__init__.py @@ -2,4 +2,4 @@ # Initialise process manager logger with Rich handler # This is the tty interface, so its designed to be coloured -get_logger("resoruce_manager", rich_handler=True) +get_logger("resource_manager", rich_handler=True) diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index 6f455423e..835fbff5e 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -73,7 +73,7 @@ def boot( # Map the requested dataflow localhost paths to realpaths, and localhost to host names for segment, _managed_objects in managed_objects.items(): - log.critical( + log.info( f"Segment '{segment}' has requested the following managed objects: {', '.join(_managed_objects)}" ) for i, managed_object in enumerate(_managed_objects): @@ -96,7 +96,7 @@ def boot( prefix = ":".join(parts[:-1]) _managed_objects[i] = f"{prefix}:{mount}" - log.critical( + log.info( f"Mapped storage path '{managed_object}' to real path '{_managed_objects[i]}'" ) From 5a038f7eddb63f82762f4bf7db07ecb313788cdf Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Fri, 8 May 2026 17:53:08 +0200 Subject: [PATCH 10/12] WIP --- src/drunc/resource_manager/client.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/drunc/resource_manager/client.py b/src/drunc/resource_manager/client.py index c1838e1f1..a1bb0aa0a 100644 --- a/src/drunc/resource_manager/client.py +++ b/src/drunc/resource_manager/client.py @@ -31,7 +31,7 @@ def _send_request(self, endpoint, payload): """ try: # verify=False is used here for local/self-signed certs (like curl -k) - self.log.critical(f"Sending request to {endpoint} with payload: {payload}") + self.log.debug(f"Sending request to {endpoint} with payload: {payload}") response = requests.post(endpoint, data=payload, verify=False) if "application/json" not in response.headers.get("Content-Type", ""): @@ -48,6 +48,7 @@ def _send_request(self, endpoint, payload): except requests.exceptions.HTTPError: self.log.warning(f"Request failed with status {response.status_code}") + self.log.debug(f"Response content: {response.text}") return response.json() except Exception as e: self.log.error(f"An unexpected error occurred: {e}") @@ -92,7 +93,7 @@ def request_resources( """ payload = { "names": ",".join(resources), - "owner": owner, + "user_name": owner, "session_id": session_id, "session_name": session_name, } @@ -100,22 +101,20 @@ def request_resources( return self._send_request(endpoint, payload) - def release_resources(self, resources: list[str], owner: str): + def release_resources(self, resources: list[str], session_id: str): """ Release resources from the Resource Manager for other runs to use. Args: resources (list[str]): List of resource names to query - owner (str): The new owner of the resources session_id (str): The session ID taking the resources - session_name (str): The session name taking the resources Returns: dict: A dictionary containing the status of the queried resources """ payload = { "names": ",".join(resources), - "owner": owner, + "session_id": session_id, } endpoint = f"{self.url}/api/release_resource/" From d06bfbbe841916eea35e971ff1a8154a3c0ce66d Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Fri, 22 May 2026 20:31:01 +0200 Subject: [PATCH 11/12] Working --- src/drunc/resource_manager/client.py | 14 ++- src/drunc/unified_shell/commands.py | 131 +++++++++++++++++++++++++-- 2 files changed, 130 insertions(+), 15 deletions(-) diff --git a/src/drunc/resource_manager/client.py b/src/drunc/resource_manager/client.py index a1bb0aa0a..a033e14df 100644 --- a/src/drunc/resource_manager/client.py +++ b/src/drunc/resource_manager/client.py @@ -45,7 +45,9 @@ def _send_request(self, endpoint, payload): # Log the successful response return response.json() - + except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e: + self.log.error(f"Failed to connect to Resource Manager at {endpoint}: {e}") + return None except requests.exceptions.HTTPError: self.log.warning(f"Request failed with status {response.status_code}") self.log.debug(f"Response content: {response.text}") @@ -56,7 +58,7 @@ def _send_request(self, endpoint, payload): def query_resources( self, resources: list[str], owner: str, session_id: str, session_name: str - ): + ) -> dict[str]: """ Query the Resource Manager for the status of the specified resources. @@ -71,14 +73,16 @@ def query_resources( """ payload = { "names": ",".join(resources), + "session_id": session_id, + "session_name": session_name, + "user_name": owner, } endpoint = f"{self.url}/api/query_resource/" - return self._send_request(endpoint, payload) def request_resources( self, resources: list[str], owner: str, session_id: str, session_name: str - ): + ) -> dict[str]: """ Request resources from the Resource Manager for isolation during the run. @@ -101,7 +105,7 @@ def request_resources( return self._send_request(endpoint, payload) - def release_resources(self, resources: list[str], session_id: str): + def release_resources(self, resources: list[str], session_id: str) -> dict[str]: """ Release resources from the Resource Manager for other runs to use. diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index 835fbff5e..7ac19c7b5 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -122,17 +122,83 @@ def boot( r for r in session_resources if not r.startswith("storage:") ] - # Request the resources from the resource manager + # Query the resources from the resource manager to check for availability if ctx.obj.resource_manager_client and ctx.obj.session_resources: log.info( - f"Requesting the following resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {', '.join(ctx.obj.session_resources)}" + f"Validating the availability of the requested resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {', '.join(ctx.obj.session_resources)}" ) - ctx.obj.resource_manager_client.request_resources( + + # Query the resource manager to check if the requested resources are available, + # and if so, request them. Note that we do this prior to booting any processes, + # to avoid booting processes and then having the resource manager deny the + # availability of the requested resources. + query_resources_response = ctx.obj.resource_manager_client.query_resources( + ctx.obj.session_resources, + getpass.getuser(), + ctx.obj.configuration_id, + session_name, + ) + + if query_resources_response.get("missing", True): + log.error( + f"The resource manager reports that the requested resources are not available. Response: {query_resources_response}" + ) + return + + # Validate that the requested resources are available in the resource manager + query_resource_response = ctx.obj.resource_manager_client.query_resources( + ctx.obj.session_resources, + getpass.getuser(), + ctx.obj.configuration_id, + session_name, + ) + unavailable_resources = [ + resource.get("name") for resource in query_resource_response.get("query_results", []) + if resource.get("session_name") != None + ] + + # If there are any unavailable resources, log them and block booting, as + # the resources required to take the run are unavailable + if unavailable_resources: + log.error(f"Resources {unavailable_resources} are not available, blocking run.") + return + else: + log.info(f"Resources {ctx.obj.session_resources} are available.") + + # Allocate the requested resources in the resource manager + request_resource_response = ctx.obj.resource_manager_client.request_resources( + ctx.obj.session_resources, + getpass.getuser(), + ctx.obj.configuration_id, + session_name, + ) + + # Check that the allocated resources match the requested resources, if not, + # log an error and block booting to avoid potential issues with processes + # booting without the required resources. Note that we check the allocated + # resources for this session and user, to avoid issues where other + # sessions/users have requested the same resources. The query checks the + # resources against both the session name and user name. + query_resource_response = ctx.obj.resource_manager_client.query_resources( ctx.obj.session_resources, getpass.getuser(), ctx.obj.configuration_id, session_name, ) + allocated_resources = [ + resource.get("name") for resource in query_resource_response.get("query_results", []) + if resource.get("session_name") == session_name and resource.get("user_name") == getpass.getuser() + ] + missing_resources = set(ctx.obj.session_resources) - set(allocated_resources) + if missing_resources: + color_coded_missing_resources_str = ", ".join([f"[red]{r.strip("'")}[/red]" for r in missing_resources]) + log.error( + f"After requesting resources, resources {color_coded_missing_resources_str} have not been allocated, stopping boot. Allocated resources will need to be manually released. " + ) + log.debug(f"Response: {request_resource_response}") + return + else: + log.info(f"Resources {ctx.obj.session_resources} have been allocated.") processes = obj.get_driver("process_manager").ps( ProcessQuery(user=user, session=session_name) @@ -255,17 +321,62 @@ def terminate(ctx, obj): log.info( f"[yellow]Empty segments (skipped):[/yellow] {', '.join(empty_segments)}" ) - ctx.obj.managed_objects = {} - ctx.obj.managed_objects_present = False - if ctx.obj.resource_manager_client and ctx.obj.session_resources: + # Query the resources from the resource manager to check for availability + if ctx.obj.resource_manager_client and ctx.obj.session_resources and ctx.obj.managed_objects_present: + released_resources_str = [f"[green]{r.strip("'")}[/]" for r in ctx.obj.session_resources] + log.info( - f"Releasing the following resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {', '.join(ctx.obj.session_resources)}" + f"Releasing the requested resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {released_resources_str}" ) - ctx.obj.resource_manager_client.release_resources( - ctx.obj.session_resources, getpass.getuser() + + # Query the resource manager to check if the requested resources are correctly + # allocated prior to releasing + query_resource_response = ctx.obj.resource_manager_client.query_resources( + ctx.obj.session_resources, + getpass.getuser(), + ctx.obj.configuration_id, + ctx.obj.session_name, ) - ctx.obj.session_resources = [] + allocated_resources = [ + resource.get("name") for resource in query_resource_response.get("query_results", []) + if resource.get("session_name") == ctx.obj.session_name and resource.get("user_name") == getpass.getuser() + ] + missing_resources = set(ctx.obj.session_resources) - set(allocated_resources) + if missing_resources: + color_coded_missing_resources_str = ", ".join([f"[red]{r.strip("'")}[/red]" for r in missing_resources]) + log.error( + f"Upon terrmination, resources {color_coded_missing_resources_str} are not allocated to session {ctx.obj.session_name}, skipping resource release. Allocated resources will need to be manually released." + ) + log.debug(f"Response: {query_resource_response}") + else: + # Release the requested resources from the resource manager + release_resource_response = ctx.obj.resource_manager_client.release_resources( + ctx.obj.session_resources, + ctx.obj.configuration_id, + ) + query_resource_response = ctx.obj.resource_manager_client.query_resources( + ctx.obj.session_resources, + getpass.getuser(), + ctx.obj.configuration_id, + ctx.obj.session_name, + ) + + # Check that the resources have been released correctly, if not, log an + # error + remaining_session_allocated_resources = [ + resource.get("name") for resource in query_resource_response.get("query_results", []) + if resource.get("session_name") == ctx.obj.session_name and resource.get("user_name") == getpass.getuser() + ] + if remaining_session_allocated_resources: + color_coded_remaining_resources_str = ", ".join([f"[red]{r.strip("'")}[/red]" for r in remaining_session_allocated_resources]) + log.critical(f"Resources {color_coded_remaining_resources_str} were not appropriately released, manually release these prior to starting any more runs.") + ctx.obj.managed_objects = {} + ctx.obj.managed_objects_present = False + else: + log.info(f"Resources {', '.join(released_resources_str)} have been released.") + ctx.obj.managed_objects = {} + ctx.obj.managed_objects_present = False obj.get_driver("process_manager").terminate() From 5d6203348e36006b6fff92f1172017bd40d06658 Mon Sep 17 00:00:00 2001 From: PawelPlesniak Date: Fri, 22 May 2026 20:36:18 +0200 Subject: [PATCH 12/12] Ruff linting --- src/drunc/unified_shell/commands.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index 7ac19c7b5..b65fb8813 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -191,7 +191,7 @@ def boot( ] missing_resources = set(ctx.obj.session_resources) - set(allocated_resources) if missing_resources: - color_coded_missing_resources_str = ", ".join([f"[red]{r.strip("'")}[/red]" for r in missing_resources]) + color_coded_missing_resources_str = ", ".join([f"[red]{r}[/red]" for r in missing_resources]) log.error( f"After requesting resources, resources {color_coded_missing_resources_str} have not been allocated, stopping boot. Allocated resources will need to be manually released. " ) @@ -324,7 +324,7 @@ def terminate(ctx, obj): # Query the resources from the resource manager to check for availability if ctx.obj.resource_manager_client and ctx.obj.session_resources and ctx.obj.managed_objects_present: - released_resources_str = [f"[green]{r.strip("'")}[/]" for r in ctx.obj.session_resources] + released_resources_str = [f"[green]{r}[/]" for r in ctx.obj.session_resources] log.info( f"Releasing the requested resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {released_resources_str}" @@ -344,14 +344,14 @@ def terminate(ctx, obj): ] missing_resources = set(ctx.obj.session_resources) - set(allocated_resources) if missing_resources: - color_coded_missing_resources_str = ", ".join([f"[red]{r.strip("'")}[/red]" for r in missing_resources]) + color_coded_missing_resources_str = ", ".join([f"[red]{r}[/red]" for r in missing_resources]) log.error( f"Upon terrmination, resources {color_coded_missing_resources_str} are not allocated to session {ctx.obj.session_name}, skipping resource release. Allocated resources will need to be manually released." ) log.debug(f"Response: {query_resource_response}") else: # Release the requested resources from the resource manager - release_resource_response = ctx.obj.resource_manager_client.release_resources( + ctx.obj.resource_manager_client.release_resources( ctx.obj.session_resources, ctx.obj.configuration_id, ) @@ -369,7 +369,7 @@ def terminate(ctx, obj): if resource.get("session_name") == ctx.obj.session_name and resource.get("user_name") == getpass.getuser() ] if remaining_session_allocated_resources: - color_coded_remaining_resources_str = ", ".join([f"[red]{r.strip("'")}[/red]" for r in remaining_session_allocated_resources]) + color_coded_remaining_resources_str = ", ".join([f"[red]{r}[/red]" for r in remaining_session_allocated_resources]) log.critical(f"Resources {color_coded_remaining_resources_str} were not appropriately released, manually release these prior to starting any more runs.") ctx.obj.managed_objects = {} ctx.obj.managed_objects_present = False