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 c29497dc8..ef9419a30 100644
--- a/config/tests/nestedConfig.data.xml
+++ b/config/tests/nestedConfig.data.xml
@@ -234,6 +234,7 @@
+
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/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..b56b7ee73
--- /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("resource_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..a033e14df
--- /dev/null
+++ b/src/drunc/resource_manager/client.py
@@ -0,0 +1,125 @@
+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.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", ""):
+ 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.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}")
+ 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
+ ) -> dict[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),
+ "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.
+
+ 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),
+ "user_name": 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], session_id: str) -> dict[str]:
+ """
+ Release resources from the Resource Manager for other runs to use.
+
+ Args:
+ resources (list[str]): List of resource names to query
+ session_id (str): The session ID taking the resources
+
+ Returns:
+ dict: A dictionary containing the status of the queried resources
+ """
+ payload = {
+ "names": ",".join(resources),
+ "session_id": session_id,
+ }
+ 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 f1c730126..b65fb8813 100644
--- a/src/drunc/unified_shell/commands.py
+++ b/src/drunc/unified_shell/commands.py
@@ -1,13 +1,18 @@
import getpass
+import os
+import socket
import sys
import click
+import conffwk
+import confmodel_dal
from druncschema.process_manager_pb2 import ProcessInstance, ProcessQuery
from drunc.controller.interface.shell_utils import controller_setup
from drunc.exceptions import DruncSetupException
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
@@ -26,14 +31,175 @@
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
+ session_resources: list[str] = []
+ segments = session_dal.segment.segments
+ while segments:
+ nested_segments = []
+ for segment in segments:
+ 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]
+ segments = nested_segments
+ 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.info(
+ 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.info(
+ 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}
+
+ # 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)}"
+ )
+
+ # Remove storage related ones for initial prototyping
+ ctx.obj.session_resources = [
+ r for r in session_resources if not r.startswith("storage:")
+ ]
+
+ # 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"Validating the availability of the requested resources from the resource manager at '{ctx.obj.resource_manager_client.url}': {', '.join(ctx.obj.session_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}[/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)
)
@@ -124,6 +290,97 @@ 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)}"
+ )
+
+ # 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}[/]" 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}"
+ )
+
+ # 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,
+ )
+ 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}[/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
+ 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}[/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()
+
+
@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 f5c9c759e..5956cab80 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
@@ -27,6 +28,12 @@ def __init__(self):
self.override_logs = True
self.running_mode = UnifiedShellMode.INTERACTIVE
self.batch_commands: list(str) = []
+ self.managed_objects: dict[
+ 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 af8bedbec..933725a2c 100644
--- a/src/drunc/unified_shell/shell.py
+++ b/src/drunc/unified_shell/shell.py
@@ -55,11 +55,11 @@
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.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
from drunc.utils.configuration import ConfTypes, OKSKey
@@ -184,6 +184,19 @@ 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:
+ ctx.obj.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}"
+ ctx.obj.log.info(
+ f"Resource manager URL parsed from configuration: [green]{resource_manager_url}[/green]"
+ )
+ ctx.obj.resource_manager_client = ResourceManagerClient(resource_manager_url)
+
ctx.obj.log.info(
f"[green]Setting up to use the process manager[/green] with configuration "
f"[green]{process_manager}[/green] and configuration id [green]"
@@ -314,8 +327,10 @@ def unified_shell(
# Add the unified shell Click commands to the CLI
ctx.obj.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
ctx.obj.log.debug("Adding [green]process_manager[/green] commands")
@@ -424,7 +439,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:
ctx.obj.log.warning(
@@ -466,7 +488,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:
+ ctx.obj.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 6f32b7aae..a11996bb7 100644
--- a/src/drunc/unified_shell/shell_utils.py
+++ b/src/drunc/unified_shell/shell_utils.py
@@ -214,3 +214,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}")