diff --git a/pyproject.toml b/pyproject.toml index 8c2d6751..16b2361c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "deadline-job-attachments == 0.1.3", # Pinned to patch version due to Host Config Script runner usage of private OpenJD Sessions API. "openjd-sessions == 0.10.14", - "openjd-model >= 0.11.3, < 0.12", + "openjd-model >= 0.11.4, < 0.12", # tomli became tomllib in standard library in Python 3.11 "tomli == 2.0.* ; python_version<'3.11'", "tomlkit >= 0.13,< 0.16", diff --git a/src/deadline_worker_agent/api_models.py b/src/deadline_worker_agent/api_models.py index e3ca5230..acecd5e1 100644 --- a/src/deadline_worker_agent/api_models.py +++ b/src/deadline_worker_agent/api_models.py @@ -218,6 +218,9 @@ class StepDetailsData(StepDetailsIdentifierFields): dependencies: NotRequired[list[str]] """A list of step identifiers that this step depends on""" + resolvedSymbolTable: NotRequired[str] + """Pre-resolved symbol table as a JSON string, forwarded to the Rust session runtime.""" + extensions: NotRequired[list[str]] """The extensions enabled for the job, as supplied by the service""" @@ -401,6 +404,8 @@ class EnvironmentDetailsData(EnvironmentDetailsIdentifierFields): """The Open Job Description schema version""" template: dict[str, Any] """The template of the environment.""" + resolvedSymbolTable: NotRequired[str] + """Pre-resolved symbol table as a JSON string, forwarded to the Rust session runtime.""" extensions: NotRequired[list[str]] """The extensions enabled for the job, as supplied by the service""" diff --git a/src/deadline_worker_agent/scheduler/scheduler.py b/src/deadline_worker_agent/scheduler/scheduler.py index 63ec8ff0..9aca6e61 100644 --- a/src/deadline_worker_agent/scheduler/scheduler.py +++ b/src/deadline_worker_agent/scheduler/scheduler.py @@ -1227,6 +1227,8 @@ def _create_new_sessions( region=self._boto_session.region_name, ) + resolved_symbol_table_json = queue.peek_resolved_symbol_table_json() + try: session = Session( id=new_session_id, @@ -1244,6 +1246,7 @@ def _create_new_sessions( session_root_dir=self._session_root_dir, farm_id=self._farm_id, region=self._boto_session.region_name, + resolved_symbol_table_json=resolved_symbol_table_json, ) except (ValueError, NotImplementedError, OSError) as e: # Runtime construction can fail per-session (e.g. the selected runtime's diff --git a/src/deadline_worker_agent/scheduler/session_queue.py b/src/deadline_worker_agent/scheduler/session_queue.py index 1c801c4b..5b9f1372 100644 --- a/src/deadline_worker_agent/scheduler/session_queue.py +++ b/src/deadline_worker_agent/scheduler/session_queue.py @@ -452,6 +452,7 @@ def dequeue(self) -> SessionActionDefinition | None: next_action = ExitEnvironmentAction( id=action_id, environment_id=environment_id, + details=environment_details, ) else: raise ValueError(f'Unknown action type "{action_type}".') @@ -557,3 +558,51 @@ def dequeue(self) -> SessionActionDefinition | None: f'Unknown action type "{action_type}". Complete action = {action_definition}' ) return next_action + + def peek_resolved_symbol_table_json(self) -> str | None: + """Inspect the first queued action's resolved symbol table without consuming it. + + This accessor is non-consuming: the queue state is not mutated, and a + subsequent ``dequeue`` call will still yield the same front action. + + Entity resolution results are cached by ``JobEntities``, so the later + ``dequeue`` issues no additional service request for the same entity. + + Returns + ------- + str | None + The ``resolved_symbol_table_json`` from the first action's entity, + or None when the queue is empty or the action type has no table. + """ + if not self._actions: + return None + + action_queue_entry = self._actions[0] + action_type = action_queue_entry.definition["actionType"] + + try: + if action_type.startswith("ENV_"): + action_queue_entry = cast(EnvironmentQueueEntry, action_queue_entry) + environment_id = action_queue_entry.definition["environmentId"] + environment_details = self._job_entities.environment_details( + environment_id=environment_id + ) + return environment_details.resolved_symbol_table_json + elif action_type == "TASK_RUN": + action_queue_entry = cast(TaskRunQueueEntry, action_queue_entry) + step_id = action_queue_entry.definition["stepId"] + step_details = self._job_entities.step_details(step_id=step_id) + return step_details.resolved_symbol_table_json + else: + return None + except Exception: + # This accessor only seeds session-scoped symbols (e.g. Job.Name), + # so a failure must not break session creation. The subsequent + # dequeue surfaces the real error through the normal action-failure + # path. + logger.warning( + "Failed to prefetch resolved symbol table for the first queued action " + "(type=%s); proceeding without it.", + action_type, + ) + return None diff --git a/src/deadline_worker_agent/sessions/actions/enter_env.py b/src/deadline_worker_agent/sessions/actions/enter_env.py index a2721baa..674491cb 100644 --- a/src/deadline_worker_agent/sessions/actions/enter_env.py +++ b/src/deadline_worker_agent/sessions/actions/enter_env.py @@ -109,4 +109,5 @@ def start( job_env_id=self._job_env_id, environment=self._details.environment, os_env_vars={"DEADLINE_SESSIONACTION_ID": self._id}, + resolved_symbol_table_json=self._details.resolved_symbol_table_json, ) diff --git a/src/deadline_worker_agent/sessions/actions/exit_env.py b/src/deadline_worker_agent/sessions/actions/exit_env.py index 739bddd8..482a2f94 100644 --- a/src/deadline_worker_agent/sessions/actions/exit_env.py +++ b/src/deadline_worker_agent/sessions/actions/exit_env.py @@ -8,6 +8,7 @@ from .openjd_action import OpenjdAction if TYPE_CHECKING: + from ..job_entities import EnvironmentDetails from ..session import Session @@ -20,26 +21,32 @@ class ExitEnvironmentAction(OpenjdAction): A unique identifier for the session action environment_id : str The job environment identifier + details : EnvironmentDetails | None + Optional environment details carrying the pre-resolved symbol table """ _environment_id: str + _details: EnvironmentDetails | None def __init__( self, *, id: str, environment_id: str, + details: EnvironmentDetails | None = None, ) -> None: super(ExitEnvironmentAction, self).__init__( id=id, action_log_kind=SessionActionLogKind.ENV_EXIT ) self._environment_id = environment_id + self._details = details def __eq__(self, other: Any) -> bool: return ( type(self) is type(other) and self._id == other._id and self._environment_id == other._environment_id + and self._details == other._details ) def start( @@ -58,5 +65,9 @@ def start( An executor for running futures """ session.exit_environment( - job_env_id=self._environment_id, os_env_vars={"DEADLINE_SESSIONACTION_ID": self._id} + job_env_id=self._environment_id, + os_env_vars={"DEADLINE_SESSIONACTION_ID": self._id}, + resolved_symbol_table_json=self._details.resolved_symbol_table_json + if self._details is not None + else None, ) diff --git a/src/deadline_worker_agent/sessions/actions/run_step_task.py b/src/deadline_worker_agent/sessions/actions/run_step_task.py index ee0ac4ee..1b1c71a4 100644 --- a/src/deadline_worker_agent/sessions/actions/run_step_task.py +++ b/src/deadline_worker_agent/sessions/actions/run_step_task.py @@ -85,4 +85,5 @@ def start(self, *, session: Session, executor: Executor) -> None: task_parameter_values=self._task_parameter_values, os_env_vars=env_vars, step_name=self._details.step_template.name, + resolved_symbol_table_json=self._details.resolved_symbol_table_json, ) diff --git a/src/deadline_worker_agent/sessions/job_entities/environment_details.py b/src/deadline_worker_agent/sessions/job_entities/environment_details.py index 7e5450aa..b50bd8e0 100644 --- a/src/deadline_worker_agent/sessions/job_entities/environment_details.py +++ b/src/deadline_worker_agent/sessions/job_entities/environment_details.py @@ -24,6 +24,9 @@ class EnvironmentDetails: environment: EnvironmentModel """The environment""" + resolved_symbol_table_json: str | None = None + """Pre-resolved symbol table JSON from the service, forwarded to the Rust session runtime.""" + @classmethod def from_boto(cls, environment_details_data: EnvironmentDetailsData) -> EnvironmentDetails: """Converts an environmentDetails entity received from BatchGetJobEntity API response into @@ -59,7 +62,10 @@ def from_boto(cls, environment_details_data: EnvironmentDetailsData) -> Environm else: raise UnsupportedSchema(schema_version.value) - return EnvironmentDetails(environment=environment) + return EnvironmentDetails( + environment=environment, + resolved_symbol_table_json=environment_details_data.get("resolvedSymbolTable", None), + ) @classmethod def validate_entity_data(cls, entity_data: dict[str, Any]) -> EnvironmentDetailsData: @@ -90,6 +96,7 @@ def validate_entity_data(cls, entity_data: dict[str, Any]) -> EnvironmentDetails Field(key="environmentId", expected_type=str, required=True), Field(key="jobId", expected_type=str, required=True), Field(key="schemaVersion", expected_type=str, required=True), + Field(key="resolvedSymbolTable", expected_type=str, required=False), Field(key="extensions", expected_type=list, required=False), ), ) diff --git a/src/deadline_worker_agent/sessions/job_entities/step_details.py b/src/deadline_worker_agent/sessions/job_entities/step_details.py index 5d4f54e2..6efe39e0 100644 --- a/src/deadline_worker_agent/sessions/job_entities/step_details.py +++ b/src/deadline_worker_agent/sessions/job_entities/step_details.py @@ -37,6 +37,9 @@ class StepDetails: dependencies: list[str] = field(default_factory=list) """The dependencies (a list of IDs) that the step depends on""" + resolved_symbol_table_json: str | None = None + """Pre-resolved symbol table JSON from the service, forwarded to the Rust session runtime.""" + @classmethod def from_boto(cls, step_details_data: StepDetailsData) -> StepDetails: """Converts an stepDetails entity received from BatchGetJobEntity API response into a @@ -86,6 +89,7 @@ def from_boto(cls, step_details_data: StepDetailsData) -> StepDetails: step_template=step_template, step_id=step_details_data["stepId"], dependencies=step_details_data["dependencies"], + resolved_symbol_table_json=step_details_data.get("resolvedSymbolTable", None), ) @classmethod @@ -118,6 +122,7 @@ def validate_entity_data(cls, entity_data: dict[str, Any]) -> StepDetailsData: Field(key="template", expected_type=dict, required=True), Field(key="stepId", expected_type=str, required=True), Field(key="dependencies", expected_type=list, required=False), + Field(key="resolvedSymbolTable", expected_type=str, required=False), Field(key="extensions", expected_type=list, required=False), ), ) diff --git a/src/deadline_worker_agent/sessions/runtime/_abc.py b/src/deadline_worker_agent/sessions/runtime/_abc.py index ccfc6ccf..bbb0c286 100644 --- a/src/deadline_worker_agent/sessions/runtime/_abc.py +++ b/src/deadline_worker_agent/sessions/runtime/_abc.py @@ -57,6 +57,7 @@ def enter_environment( environment: EnvironmentModel, identifier: Optional[EnvironmentIdentifier] = None, os_env_vars: Optional[dict[str, str]] = None, + resolved_symbol_table_json: str | None = None, ) -> EnvironmentIdentifier: """Enter an environment; returns its identifier.""" ... @@ -68,6 +69,7 @@ def exit_environment( identifier: EnvironmentIdentifier, os_env_vars: Optional[dict[str, str]] = None, keep_session_running: bool = False, + resolved_symbol_table_json: str | None = None, ) -> None: """Exit a previously entered environment.""" ... @@ -81,6 +83,7 @@ def run_task( os_env_vars: Optional[dict[str, str]] = None, log_task_banner: bool = True, step_name: str | None = None, + resolved_symbol_table_json: str | None = None, ) -> None: """Run a task within the session's active environment(s).""" ... diff --git a/src/deadline_worker_agent/sessions/runtime/_config.py b/src/deadline_worker_agent/sessions/runtime/_config.py index bbba6339..7bc71931 100644 --- a/src/deadline_worker_agent/sessions/runtime/_config.py +++ b/src/deadline_worker_agent/sessions/runtime/_config.py @@ -37,3 +37,9 @@ class SessionRuntimeConfig: session_root_directory: Path spec_revision: str = "2023-09" supported_extensions: tuple[str, ...] = () + resolved_symbol_table_json: str | None = None + """Pre-resolved symbol table JSON observed on the session's first queued action. + + Adapters that seed session-scoped symbols at construction (e.g. job_name for + the classic Python session) consume this at build time. + """ diff --git a/src/deadline_worker_agent/sessions/runtime/python.py b/src/deadline_worker_agent/sessions/runtime/python.py index ddade84c..affce660 100644 --- a/src/deadline_worker_agent/sessions/runtime/python.py +++ b/src/deadline_worker_agent/sessions/runtime/python.py @@ -3,9 +3,11 @@ from __future__ import annotations from datetime import timedelta +from logging import getLogger from pathlib import Path from typing import TYPE_CHECKING, Any, Optional +from openjd.expr import SerializedSymbolTable from openjd.model import RevisionExtensions, SpecificationRevision from openjd.sessions import Session as OpenJDSession @@ -22,6 +24,29 @@ __all__ = ["PythonSessionRuntime"] +logger = getLogger(__name__) + + +def _extract_job_name(json_str: str | None) -> str | None: + """Extract the Job.Name value from a resolved symbol table JSON string. + + Returns None when the input is None, the table lacks a Job.Name entry, the + entry's value is not a string, or parsing fails (graceful degradation — + mirrors _parse_resolved_symtab in the Rust adapter). + """ + if json_str is None: + return None + try: + symtab = SerializedSymbolTable.from_json_str(json_str).to_symtab() + entry = symtab.get("Job.Name") + if entry is None: + return None + value = entry.item() + return value if isinstance(value, str) else None + except Exception as e: + logger.warning("Failed to extract Job.Name from resolvedSymbolTable: %s", e) + return None + class PythonSessionRuntime(SessionRuntime): """SessionRuntime backed by openjd.sessions (v0 Python implementation).""" @@ -38,6 +63,7 @@ def __init__(self, config: SessionRuntimeConfig) -> None: callback=config.action_callback, os_env_vars=config.os_env_vars, session_root_directory=config.session_root_directory, + job_name=_extract_job_name(config.resolved_symbol_table_json), revision_extensions=RevisionExtensions( # Currently for simplicity request that our session allow all extensions. # This does not obey the spec. It should be changed at a later date to the @@ -53,7 +79,10 @@ def enter_environment( environment: EnvironmentModel, identifier: Optional[EnvironmentIdentifier] = None, os_env_vars: Optional[dict[str, str]] = None, + resolved_symbol_table_json: str | None = None, ) -> EnvironmentIdentifier: + # resolved_symbol_table_json: not forwarded — the v0 Python session does + # not support pre-resolved symbol tables. return self._session.enter_environment( environment=environment, identifier=identifier, @@ -66,7 +95,10 @@ def exit_environment( identifier: EnvironmentIdentifier, os_env_vars: Optional[dict[str, str]] = None, keep_session_running: bool = False, + resolved_symbol_table_json: str | None = None, ) -> None: + # resolved_symbol_table_json: not forwarded — the v0 Python session does + # not support pre-resolved symbol tables. self._session.exit_environment( identifier=identifier, os_env_vars=os_env_vars, @@ -81,7 +113,10 @@ def run_task( os_env_vars: Optional[dict[str, str]] = None, log_task_banner: bool = True, step_name: str | None = None, + resolved_symbol_table_json: str | None = None, ) -> None: + # resolved_symbol_table_json: not forwarded — the v0 Python session does + # not support pre-resolved symbol tables. self._session.run_task( step_script=step_script, task_parameter_values=task_parameter_values, diff --git a/src/deadline_worker_agent/sessions/runtime/rust.py b/src/deadline_worker_agent/sessions/runtime/rust.py index 8ec36ed5..84edd92d 100644 --- a/src/deadline_worker_agent/sessions/runtime/rust.py +++ b/src/deadline_worker_agent/sessions/runtime/rust.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Optional from openjd._openjd_rs import create_environment, deserialize_step -from openjd.expr import PathFormat, PathMappingRule as RustPathMappingRule +from openjd.expr import PathFormat, PathMappingRule as RustPathMappingRule, SerializedSymbolTable from openjd.model._v1 import decode_environment_template from openjd.model._v1.types import ( JobParameterType, @@ -269,6 +269,21 @@ def _to_rust_model_extensions( return extensions, converted_names +def _parse_resolved_symtab(json_str: str | None) -> Any: + """Parse a resolved symbol table JSON string into a SerializedSymbolTable. + + Returns None when the input is None or when parsing fails (graceful + degradation — the session proceeds without the pre-resolved table). + """ + if json_str is None: + return None + try: + return SerializedSymbolTable.from_json_str(json_str) + except Exception as e: + logger.warning("Failed to parse resolvedSymbolTable; proceeding without it: %s", e) + return None + + class RustSessionRuntime(SessionRuntime): """SessionRuntime backed by openjd.sessions._v1 (Rust implementation).""" @@ -345,6 +360,7 @@ def enter_environment( environment: EnvironmentModel, identifier: Optional[EnvironmentIdentifier] = None, os_env_vars: Optional[dict[str, str]] = None, + resolved_symbol_table_json: str | None = None, ) -> EnvironmentIdentifier: # The shared action layer hands a pydantic v2023_09 environment, but the # Rust session needs a native _v1 environment. Serialize to the OpenJD @@ -372,10 +388,15 @@ def enter_environment( supported_extensions=list(self._decode_extensions) or None, ) ) + + # Parse the pre-resolved symbol table if the service provided one. + resolved_symtab = _parse_resolved_symtab(resolved_symbol_table_json) + return self._session.enter_environment( environment=native_environment, identifier=identifier, os_env_vars=os_env_vars, + resolved_symtab=resolved_symtab, ) @convert_runtime_crashes @@ -385,11 +406,14 @@ def exit_environment( identifier: EnvironmentIdentifier, os_env_vars: Optional[dict[str, str]] = None, keep_session_running: bool = False, + resolved_symbol_table_json: str | None = None, ) -> None: + resolved_symtab = _parse_resolved_symtab(resolved_symbol_table_json) self._session.exit_environment( identifier=identifier, os_env_vars=os_env_vars, keep_session_running=keep_session_running, + resolved_symtab=resolved_symtab, ) @convert_runtime_crashes @@ -401,6 +425,7 @@ def run_task( os_env_vars: Optional[dict[str, str]] = None, log_task_banner: bool = True, step_name: str | None = None, + resolved_symbol_table_json: str | None = None, ) -> None: # step_name: forwarded to the _v1 session so RFC 0008's WrappedStep.Name # resolves correctly inside onWrapTaskRun hooks. @@ -419,12 +444,17 @@ def run_task( "script": step_script.model_dump(mode="json", by_alias=True, exclude_none=True), } ).script + + # Parse the pre-resolved symbol table if the service provided one. + resolved_symtab = _parse_resolved_symtab(resolved_symbol_table_json) + self._session.run_task( step_script=native_step_script, task_parameter_values=_to_rust_task_parameter_values(task_parameter_values), os_env_vars=os_env_vars, log_task_banner=log_task_banner, step_name=step_name, + resolved_symtab=resolved_symtab, ) @convert_runtime_crashes diff --git a/src/deadline_worker_agent/sessions/session.py b/src/deadline_worker_agent/sessions/session.py index 051d58cb..d13235f0 100644 --- a/src/deadline_worker_agent/sessions/session.py +++ b/src/deadline_worker_agent/sessions/session.py @@ -101,6 +101,9 @@ class ActiveEnvironment: job_env_id: str """A unique identifier that identifies the environment within the Open Job Description job model""" + resolved_symbol_table_json: str | None = None + """Retained so environment teardown resolves the same symbols the enter action used.""" + @dataclass(frozen=True) class CurrentAction: @@ -180,6 +183,7 @@ def __init__( session_root_dir: Path, farm_id: str = "", region: Optional[str] = None, + resolved_symbol_table_json: str | None = None, ) -> None: self._id = id self._session_runtime_kind = session_runtime_kind @@ -221,6 +225,7 @@ def openjd_session_action_callback(session_id: str, action_status: ActionStatus) session_root_directory=session_root_dir, spec_revision="2023-09", supported_extensions=session_extensions(self._job_details.extensions), + resolved_symbol_table_json=resolved_symbol_table_json, ), ) @@ -465,7 +470,11 @@ def _cleanup(self) -> None: # After canceling the running action, we exit any active environments actions.extend( ( - partial(self._runtime.exit_environment, identifier=env.session_env_id), + partial( + self._runtime.exit_environment, + identifier=env.session_env_id, + resolved_symbol_table_json=env.resolved_symbol_table_json, + ), f"exit environment {env.job_env_id}", ) for env in reversed(self._active_envs) @@ -869,14 +878,19 @@ def enter_environment( job_env_id: str, environment: EnvironmentModel, os_env_vars: Optional[dict[str, str]] = None, + resolved_symbol_table_json: str | None = None, ) -> None: session_env_id = self._runtime.enter_environment( - environment=environment, identifier=job_env_id, os_env_vars=os_env_vars + environment=environment, + identifier=job_env_id, + os_env_vars=os_env_vars, + resolved_symbol_table_json=resolved_symbol_table_json, ) self._active_envs.append( ActiveEnvironment( job_env_id=job_env_id, session_env_id=session_env_id, + resolved_symbol_table_json=resolved_symbol_table_json, ) ) @@ -885,6 +899,7 @@ def exit_environment( *, job_env_id: str, os_env_vars: Optional[dict[str, str]] = None, + resolved_symbol_table_json: str | None = None, ) -> None: if not self._active_envs or self._active_envs[-1].job_env_id != job_env_id: env_stack_str = ", ".join(env.job_env_id for env in self._active_envs) @@ -893,8 +908,18 @@ def exit_environment( f"Active environments from outer-most to inner-most are: {env_stack_str}" ) active_env = self._active_envs[-1] + # The exit action's own table is preferred; the enter-time table is used + # when the service omits one, because the runtime does not replay + # enter-time symbols. + effective_table = ( + resolved_symbol_table_json + if resolved_symbol_table_json is not None + else active_env.resolved_symbol_table_json + ) self._runtime.exit_environment( - identifier=active_env.session_env_id, os_env_vars=os_env_vars + identifier=active_env.session_env_id, + os_env_vars=os_env_vars, + resolved_symbol_table_json=effective_table, ) self._active_envs.pop() @@ -1202,6 +1227,7 @@ def run_task( os_env_vars: Optional[dict[str, str]] = None, log_task_banner: bool = True, step_name: str | None = None, + resolved_symbol_table_json: str | None = None, ) -> None: self._runtime.run_task( step_script=step_script, @@ -1209,6 +1235,7 @@ def run_task( os_env_vars=os_env_vars, log_task_banner=log_task_banner, step_name=step_name, + resolved_symbol_table_json=resolved_symbol_table_json, ) def _run_attachment_sync_task( diff --git a/test/unit/scheduler/test_scheduler.py b/test/unit/scheduler/test_scheduler.py index dcce4608..630bfd29 100644 --- a/test/unit/scheduler/test_scheduler.py +++ b/test/unit/scheduler/test_scheduler.py @@ -1480,6 +1480,79 @@ def test_determine_user_for_session( assert result == expected_result +class TestCreateNewSessionsPrefetchSymbolTable: + """Tests that the scheduler prefetches the resolved symbol table and passes it to Session""" + + @pytest.fixture + def mock_job_entities(self) -> Generator[MagicMock, None, None]: + with patch.object(scheduler_mod, "JobEntities") as job_entities_mock: + yield job_entities_mock + + def test_passes_prefetched_symbol_table_to_session( + self, + scheduler: WorkerScheduler, + mock_session: MagicMock, + mock_job_entities: MagicMock, + ) -> None: + """Tests that the scheduler calls peek_resolved_symbol_table_json on the queue + and passes the result to Session(...) as resolved_symbol_table_json.""" + # GIVEN + table_json = '[{"name":"Job.Name","type":"string","value":"Example Job"}]' + queue_id = "queue-abcdef0123456789abcdef0123456789" + session_id = "session-abcdef0123456789abcdef0123456789" + scheduler._job_run_as_user_override = JobsRunAsUserOverride(run_as_agent=False) + assigned_sessions: dict[str, AssignedSession] = { + session_id: AssignedSession( + queueId=queue_id, + jobId="job-abcdef0123456789abcdef0123456789", + logConfiguration=LogConfiguration( + logDriver="awslogs", + options={}, + parameters={"interval": "15"}, + ), + sessionActions=[ + EnvironmentAction( + actionType="ENV_ENTER", + environmentId="env-1", + sessionActionId="action-1", + ), + ], + ), + } + job_entity_mock = MagicMock() + job_entity_mock.job_details.return_value = JobDetails( + log_group_name="/aws/deadline/queue-0000", + schema_version=SpecificationRevision.v2023_09, + job_run_as_user=JobRunAsUser( + posix=( + PosixSessionUser(user="username", group="group") if os.name == "posix" else None + ), + windows=( + WindowsSessionUser(user="username", password="password") + if os.name == "nt" + else None + ), + windows_settings=None, + ), + ) + mock_job_entities.return_value = job_entity_mock + + with ( + patch.object(scheduler, "_executor"), + patch.object( + scheduler_mod.SessionActionQueue, + "peek_resolved_symbol_table_json", + return_value=table_json, + ), + ): + # WHEN + scheduler._create_new_sessions(assigned_sessions=assigned_sessions) + + # THEN + mock_session.assert_called_once() + assert mock_session.call_args.kwargs["resolved_symbol_table_json"] == table_json + + class TestCreateNewSessionsRuntimeHint: """Tests for runtime hint consumption in WorkerScheduler._create_new_sessions""" diff --git a/test/unit/scheduler/test_session_queue.py b/test/unit/scheduler/test_session_queue.py index 0493645c..47f5232d 100644 --- a/test/unit/scheduler/test_session_queue.py +++ b/test/unit/scheduler/test_session_queue.py @@ -31,6 +31,7 @@ AttachmentDownloadActionQueueEntry, AttachmentUploadActionQueueEntry, ) +import deadline_worker_agent.scheduler.session_queue as session_queue_mod from deadline_worker_agent.sessions.actions import ( EnterEnvironmentAction, @@ -131,6 +132,9 @@ class TestSessionActionQueueDequeue: ExitEnvironmentAction( id="id", environment_id="envid", + details=EnvironmentDetails( + environment=Environment(name="TestEnv", script=_TEST_ENVIRONMENT_SCRIPT) + ), ), id="env exit", ), @@ -586,3 +590,198 @@ def test_list_all_action_identifiers( # THEN assert identifiers == expected_identifiers + + +class TestPeekResolvedSymbolTableJson: + """Tests for SessionActionQueue.peek_resolved_symbol_table_json""" + + def test_returns_none_for_empty_queue( + self, + session_queue: SessionActionQueue, + ) -> None: + # GIVEN + assert session_queue._actions == [] + + # WHEN + result = session_queue.peek_resolved_symbol_table_json() + + # THEN + assert result is None + + def test_returns_environment_table_when_first_action_is_env_enter( + self, + session_queue: SessionActionQueue, + job_entities: MagicMock, + ) -> None: + # GIVEN + table_json = '[{"name":"Job.Name","type":"string","value":"Example Job"}]' + job_entities.environment_details.return_value = EnvironmentDetails( + environment=Environment(name="TestEnv", script=_TEST_ENVIRONMENT_SCRIPT), + resolved_symbol_table_json=table_json, + ) + entry = EnvironmentQueueEntry( + Mock(), + EnvironmentAction( + sessionActionId="action-1", actionType="ENV_ENTER", environmentId="env-1" + ), + ) + session_queue._actions = [entry] + session_queue._actions_by_id["action-1"] = entry + + # WHEN + result = session_queue.peek_resolved_symbol_table_json() + + # THEN + assert result == table_json + job_entities.environment_details.assert_called_once_with(environment_id="env-1") + + def test_returns_environment_table_when_first_action_is_env_exit( + self, + session_queue: SessionActionQueue, + job_entities: MagicMock, + ) -> None: + # GIVEN + table_json = '[{"name":"Job.Name","type":"string","value":"Example Job"}]' + job_entities.environment_details.return_value = EnvironmentDetails( + environment=Environment(name="TestEnv", script=_TEST_ENVIRONMENT_SCRIPT), + resolved_symbol_table_json=table_json, + ) + entry = EnvironmentQueueEntry( + Mock(), + EnvironmentAction( + sessionActionId="action-1", actionType="ENV_EXIT", environmentId="env-1" + ), + ) + session_queue._actions = [entry] + session_queue._actions_by_id["action-1"] = entry + + # WHEN + result = session_queue.peek_resolved_symbol_table_json() + + # THEN + assert result == table_json + job_entities.environment_details.assert_called_once_with(environment_id="env-1") + + def test_returns_step_table_when_first_action_is_task_run( + self, + session_queue: SessionActionQueue, + job_entities: MagicMock, + ) -> None: + # GIVEN + table_json = '[{"name":"Job.Name","type":"string","value":"Example Job"}]' + job_entities.step_details.return_value = StepDetails( + step_template=_TEST_STEP_TEMPLATE, + step_id="step-1", + resolved_symbol_table_json=table_json, + ) + entry = TaskRunQueueEntry( + Mock(), + TaskRunAction( + sessionActionId="action-1", + actionType="TASK_RUN", + stepId="step-1", + taskId="task-1", + parameters={}, + ), + ) + session_queue._actions = [entry] + session_queue._actions_by_id["action-1"] = entry + + # WHEN + result = session_queue.peek_resolved_symbol_table_json() + + # THEN + assert result == table_json + job_entities.step_details.assert_called_once_with(step_id="step-1") + + def test_returns_none_and_does_not_raise_when_entity_resolution_raises( + self, + session_queue: SessionActionQueue, + job_entities: MagicMock, + ) -> None: + # GIVEN + job_entities.environment_details.side_effect = RuntimeError("service unavailable") + entry = EnvironmentQueueEntry( + Mock(), + EnvironmentAction( + sessionActionId="action-1", actionType="ENV_ENTER", environmentId="env-1" + ), + ) + session_queue._actions = [entry] + session_queue._actions_by_id["action-1"] = entry + + # WHEN + with patch.object(session_queue_mod, "logger") as mock_logger: + result = session_queue.peek_resolved_symbol_table_json() + + # THEN + assert result is None + mock_logger.warning.assert_called_once() + + def test_does_not_consume_queue( + self, + session_queue: SessionActionQueue, + job_entities: MagicMock, + ) -> None: + # GIVEN + table_json = '[{"name":"Job.Name","type":"string","value":"Example Job"}]' + job_entities.environment_details.return_value = EnvironmentDetails( + environment=Environment(name="TestEnv", script=_TEST_ENVIRONMENT_SCRIPT), + resolved_symbol_table_json=table_json, + ) + entry = EnvironmentQueueEntry( + Mock(), + EnvironmentAction( + sessionActionId="action-1", actionType="ENV_ENTER", environmentId="env-1" + ), + ) + session_queue._actions = [entry] + session_queue._actions_by_id["action-1"] = entry + + # WHEN + peek_result = session_queue.peek_resolved_symbol_table_json() + + # THEN — queue is unmodified + assert len(session_queue._actions) == 1 + assert "action-1" in session_queue._actions_by_id + + # AND — subsequent dequeue yields the same action + dequeue_result = session_queue.dequeue() + assert dequeue_result is not None + assert dequeue_result.id == "action-1" + assert peek_result == table_json + + +class TestDequeueEnvExitDetails: + """Tests that ENV_EXIT dequeue produces an ExitEnvironmentAction carrying fetched details""" + + def test_env_exit_dequeue_includes_details( + self, + session_queue: SessionActionQueue, + job_entities: MagicMock, + ) -> None: + # GIVEN + table_json = '[{"name":"Job.Name","type":"string","value":"Example Job"}]' + env_details = EnvironmentDetails( + environment=Environment(name="TestEnv", script=_TEST_ENVIRONMENT_SCRIPT), + resolved_symbol_table_json=table_json, + ) + job_entities.environment_details.return_value = env_details + entry = EnvironmentQueueEntry( + Mock(), + EnvironmentAction( + sessionActionId="action-1", actionType="ENV_EXIT", environmentId="env-1" + ), + ) + session_queue._actions = [entry] + session_queue._actions_by_id["action-1"] = entry + + # WHEN + result = session_queue.dequeue() + + # THEN + assert isinstance(result, ExitEnvironmentAction) + assert result.id == "action-1" + # The details kwarg is passed to ExitEnvironmentAction; verify the + # attribute is set (concurrent agent adds _details field). + assert result._details is env_details # type: ignore[attr-defined] diff --git a/test/unit/sessions/actions/test_exit_env.py b/test/unit/sessions/actions/test_exit_env.py new file mode 100644 index 00000000..0c744297 --- /dev/null +++ b/test/unit/sessions/actions/test_exit_env.py @@ -0,0 +1,98 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest + +from deadline_worker_agent.sessions.actions.exit_env import ExitEnvironmentAction +from deadline_worker_agent.sessions.job_entities import EnvironmentDetails + + +@pytest.fixture +def mock_session() -> Mock: + session = Mock() + session.exit_environment = Mock() + return session + + +@pytest.fixture +def mock_executor() -> Mock: + return Mock() + + +class TestExitEnvironmentAction: + """Tests for ExitEnvironmentAction with optional details.""" + + def test_start_forwards_resolved_symbol_table_json_from_details( + self, mock_session: Mock, mock_executor: Mock + ) -> None: + """start() forwards self._details.resolved_symbol_table_json to session.exit_environment.""" + table = '[{"name":"Job.Name","type":"string","value":"Outer"}]' + details = Mock(spec=EnvironmentDetails) + details.resolved_symbol_table_json = table + + action = ExitEnvironmentAction( + id="action-exit-1", + environment_id="env-outer", + details=details, + ) + + action.start(session=mock_session, executor=mock_executor) + + mock_session.exit_environment.assert_called_once_with( + job_env_id="env-outer", + os_env_vars={"DEADLINE_SESSIONACTION_ID": "action-exit-1"}, + resolved_symbol_table_json=table, + ) + + def test_start_forwards_none_when_no_details( + self, mock_session: Mock, mock_executor: Mock + ) -> None: + """start() forwards None when constructed without details.""" + action = ExitEnvironmentAction( + id="action-exit-2", + environment_id="env-inner", + ) + + action.start(session=mock_session, executor=mock_executor) + + mock_session.exit_environment.assert_called_once_with( + job_env_id="env-inner", + os_env_vars={"DEADLINE_SESSIONACTION_ID": "action-exit-2"}, + resolved_symbol_table_json=None, + ) + + def test_eq_includes_details(self) -> None: + """__eq__ considers details so equality proves the details were threaded.""" + details_a = Mock(spec=EnvironmentDetails) + details_a.resolved_symbol_table_json = '[{"name":"Job.Name","type":"string","value":"A"}]' + details_b = Mock(spec=EnvironmentDetails) + details_b.resolved_symbol_table_json = '[{"name":"Job.Name","type":"string","value":"B"}]' + + action_a = ExitEnvironmentAction(id="action-1", environment_id="env-1", details=details_a) + action_b = ExitEnvironmentAction(id="action-1", environment_id="env-1", details=details_b) + action_c = ExitEnvironmentAction(id="action-1", environment_id="env-1", details=details_a) + + assert action_a == action_c + assert action_a != action_b + + def test_eq_without_details(self) -> None: + """Two actions without details are equal when id and environment_id match.""" + action_a = ExitEnvironmentAction(id="action-1", environment_id="env-1") + action_b = ExitEnvironmentAction(id="action-1", environment_id="env-1") + + assert action_a == action_b + + def test_eq_none_details_vs_details(self) -> None: + """An action without details is not equal to one with details.""" + details = Mock(spec=EnvironmentDetails) + details.resolved_symbol_table_json = "[]" + + action_no_details = ExitEnvironmentAction(id="action-1", environment_id="env-1") + action_with_details = ExitEnvironmentAction( + id="action-1", environment_id="env-1", details=details + ) + + assert action_no_details != action_with_details diff --git a/test/unit/sessions/job_entities/test_environment_details.py b/test/unit/sessions/job_entities/test_environment_details.py index 7dd70f74..c964ad4e 100644 --- a/test/unit/sessions/job_entities/test_environment_details.py +++ b/test/unit/sessions/job_entities/test_environment_details.py @@ -151,3 +151,31 @@ def test_from_boto_parses_wrap_actions_environment(self) -> None: assert actions.onWrapEnvExit is not None assert actions.onWrapEnvExit.command == "/bin/echo" assert actions.onWrapEnvExit.args == ["exit"] + + +class TestResolvedSymbolTable: + """Tests for the resolvedSymbolTable field on EnvironmentDetails.""" + + def test_from_boto_extracts_resolved_symbol_table_when_present(self) -> None: + """from_boto sets resolved_symbol_table_json when the field is present.""" + symtab_json = '[{"name":"Job.Name","type":"string","value":"MyJob"}]' + environment_details_data = { + "jobId": "job-0000", + "environmentId": "env-0000", + "schemaVersion": "jobtemplate-2023-09", + "template": { + "name": "TestEnv", + "script": { + "actions": { + "onEnter": {"command": "/bin/echo", "args": ["enter"]}, + } + }, + }, + "resolvedSymbolTable": symtab_json, + } + + result = EnvironmentDetails.from_boto( + cast(EnvironmentDetailsData, environment_details_data) + ) + + assert result.resolved_symbol_table_json == symtab_json diff --git a/test/unit/sessions/job_entities/test_step_details.py b/test/unit/sessions/job_entities/test_step_details.py index 833eb660..328a9617 100644 --- a/test/unit/sessions/job_entities/test_step_details.py +++ b/test/unit/sessions/job_entities/test_step_details.py @@ -212,3 +212,30 @@ def test_absent_extensions_falls_back(self) -> None: result = StepDetails.from_boto(cast(StepDetailsData, step_details_data)) assert result.step_template.name == "TestStep" + + +class TestResolvedSymbolTable: + """Tests for the resolvedSymbolTable field on StepDetails.""" + + def test_from_boto_extracts_resolved_symbol_table_when_present(self) -> None: + """from_boto sets resolved_symbol_table_json when the field is present.""" + symtab_json = '[{"name":"Job.Name","type":"string","value":"MyJob"}]' + step_details_data = { + "jobId": "job-0000", + "schemaVersion": "jobtemplate-2023-09", + "stepId": "step-0000", + "dependencies": [], + "template": { + "name": "TestStep", + "script": { + "actions": { + "onRun": {"command": "/bin/echo", "args": ["hello"]}, + } + }, + }, + "resolvedSymbolTable": symtab_json, + } + + result = StepDetails.from_boto(cast(StepDetailsData, step_details_data)) + + assert result.resolved_symbol_table_json == symtab_json diff --git a/test/unit/sessions/runtime/conftest.py b/test/unit/sessions/runtime/conftest.py index 1eddb93e..390f4011 100644 --- a/test/unit/sessions/runtime/conftest.py +++ b/test/unit/sessions/runtime/conftest.py @@ -39,6 +39,7 @@ def enter_environment( environment: Any = None, identifier: Any = None, os_env_vars: Optional[dict[str, str]] = None, + resolved_symbol_table_json: str | None = None, ) -> str: return "env-id" @@ -48,6 +49,7 @@ def exit_environment( identifier: Any = None, os_env_vars: Optional[dict[str, str]] = None, keep_session_running: bool = False, + resolved_symbol_table_json: str | None = None, ) -> None: return None @@ -59,6 +61,7 @@ def run_task( os_env_vars: Optional[dict[str, str]] = None, log_task_banner: bool = True, step_name: str | None = None, + resolved_symbol_table_json: str | None = None, ) -> None: return None diff --git a/test/unit/sessions/runtime/test_python.py b/test/unit/sessions/runtime/test_python.py index 1b5c308d..a100ca05 100644 --- a/test/unit/sessions/runtime/test_python.py +++ b/test/unit/sessions/runtime/test_python.py @@ -33,7 +33,7 @@ def runtime_config() -> SessionRuntimeConfig: @pytest.fixture() def mock_openjd_session() -> Generator[MagicMock, None, None]: - with patch.object(python_module, "OpenJDSession") as mock_cls: + with patch.object(python_module, "OpenJDSession", autospec=True) as mock_cls: yield mock_cls @@ -53,6 +53,7 @@ def test_construction_when_default_config_delegates_to_openjd_session( assert call_kwargs["callback"] is runtime_config.action_callback assert call_kwargs["os_env_vars"] is None assert call_kwargs["session_root_directory"] == Path("/tmp/sessions/session-1") + assert call_kwargs["job_name"] is None rev_ext = call_kwargs["revision_extensions"] assert rev_ext.spec_rev == SpecificationRevision.v2023_09 assert rev_ext.extensions == set() @@ -123,6 +124,21 @@ def test_exit_environment_when_called_delegates_to_wrapped_session( identifier=identifier, os_env_vars={"A": "B"}, keep_session_running=True ) + def test_exit_environment_does_not_forward_resolved_symbol_table_json( + self, adapter: PythonSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """The v0 Python session has no resolved-table API; the kwarg is accepted but not forwarded.""" + identifier = MagicMock() + + adapter.exit_environment( + identifier=identifier, + resolved_symbol_table_json='[{"name":"Job.Name","type":"string","value":"X"}]', + ) + + mock_session_instance.exit_environment.assert_called_once_with( + identifier=identifier, os_env_vars=None, keep_session_running=False + ) + def test_run_task_when_called_delegates_to_wrapped_session( self, adapter: PythonSessionRuntime, mock_session_instance: MagicMock ) -> None: @@ -215,3 +231,89 @@ def test_action_status_when_none_returns_none( mock_session_instance.action_status = None assert adapter.action_status is None + + +class TestPythonSessionRuntimeJobName: + """Tests for job_name extraction and forwarding at construction time.""" + + def test_construction_passes_job_name_when_resolved_table_has_job_name( + self, mock_openjd_session: MagicMock + ) -> None: + config = SessionRuntimeConfig( + session_id="session-jn-1", + job_parameter_values={}, + path_mapping_rules=None, + retain_working_dir=False, + user=None, + action_callback=lambda session_id, status: None, + os_env_vars=None, + session_root_directory=Path("/tmp/sessions/session-jn-1"), + resolved_symbol_table_json='[{"name":"Job.Name","type":"string","value":"Example Job"}]', + ) + + PythonSessionRuntime(config) + + call_kwargs = mock_openjd_session.call_args.kwargs + assert call_kwargs["job_name"] == "Example Job" + + def test_construction_passes_none_when_resolved_table_json_is_none( + self, mock_openjd_session: MagicMock + ) -> None: + config = SessionRuntimeConfig( + session_id="session-jn-2", + job_parameter_values={}, + path_mapping_rules=None, + retain_working_dir=False, + user=None, + action_callback=lambda session_id, status: None, + os_env_vars=None, + session_root_directory=Path("/tmp/sessions/session-jn-2"), + resolved_symbol_table_json=None, + ) + + PythonSessionRuntime(config) + + call_kwargs = mock_openjd_session.call_args.kwargs + assert call_kwargs["job_name"] is None + + def test_construction_passes_none_when_resolved_table_lacks_job_name( + self, mock_openjd_session: MagicMock + ) -> None: + config = SessionRuntimeConfig( + session_id="session-jn-3", + job_parameter_values={}, + path_mapping_rules=None, + retain_working_dir=False, + user=None, + action_callback=lambda session_id, status: None, + os_env_vars=None, + session_root_directory=Path("/tmp/sessions/session-jn-3"), + resolved_symbol_table_json="[]", + ) + + PythonSessionRuntime(config) + + call_kwargs = mock_openjd_session.call_args.kwargs + assert call_kwargs["job_name"] is None + + def test_construction_passes_none_when_resolved_table_is_malformed_json( + self, mock_openjd_session: MagicMock + ) -> None: + config = SessionRuntimeConfig( + session_id="session-jn-4", + job_parameter_values={}, + path_mapping_rules=None, + retain_working_dir=False, + user=None, + action_callback=lambda session_id, status: None, + os_env_vars=None, + session_root_directory=Path("/tmp/sessions/session-jn-4"), + resolved_symbol_table_json="{not json", + ) + + with patch.object(python_module, "logger") as mock_logger: + PythonSessionRuntime(config) + + mock_logger.warning.assert_called_once() + call_kwargs = mock_openjd_session.call_args.kwargs + assert call_kwargs["job_name"] is None diff --git a/test/unit/sessions/runtime/test_rust.py b/test/unit/sessions/runtime/test_rust.py index d632b8c5..43dd764a 100644 --- a/test/unit/sessions/runtime/test_rust.py +++ b/test/unit/sessions/runtime/test_rust.py @@ -53,7 +53,7 @@ def runtime_config() -> SessionRuntimeConfig: @pytest.fixture() def mock_rust_session() -> Generator[MagicMock, None, None]: - with patch.object(rust_module, "OpenJDRustSession") as mock_cls: + with patch.object(rust_module, "OpenJDRustSession", autospec=True) as mock_cls: yield mock_cls @@ -275,6 +275,7 @@ def test_enter_environment_when_called_converts_env_and_delegates( environment=mock_create.return_value, identifier=identifier, os_env_vars=os_env, + resolved_symtab=None, ) assert result is mock_session_instance.enter_environment.return_value @@ -465,7 +466,10 @@ def test_exit_environment_when_called_delegates_to_wrapped_session( ) mock_session_instance.exit_environment.assert_called_once_with( - identifier="job-env-1", os_env_vars={"A": "B"}, keep_session_running=True + identifier="job-env-1", + os_env_vars={"A": "B"}, + keep_session_running=True, + resolved_symtab=None, ) def test_run_task_when_called_converts_step_script_and_delegates( @@ -1319,3 +1323,148 @@ def test_non_primitive_types_are_declared(self) -> None: {"name": "Flag", "type": "BOOL"}, {"name": "Expr", "type": "RANGE_EXPR"}, ] + + +class TestResolvedSymbolTableForwarding: + """Tests for resolved_symbol_table_json parsing and forwarding to the _v1 session.""" + + @pytest.fixture() + def adapter( + self, runtime_config: SessionRuntimeConfig, mock_rust_session: MagicMock + ) -> RustSessionRuntime: + return RustSessionRuntime(runtime_config) + + @pytest.fixture() + def mock_session_instance(self, mock_rust_session: MagicMock) -> MagicMock: + return mock_rust_session.return_value + + def test_enter_environment_forwards_resolved_symtab_when_json_present( + self, adapter: RustSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """When resolved_symbol_table_json is provided, it's parsed and forwarded.""" + environment = MagicMock() + symtab_json = '[{"name":"Job.Name","type":"string","value":"TestJob"}]' + fake_symtab = MagicMock() + + with ( + patch.object(rust_module, "decode_environment_template"), + patch.object(rust_module, "create_environment"), + patch.object( + rust_module.SerializedSymbolTable, "from_json_str", return_value=fake_symtab + ) as mock_from_json, + ): + adapter.enter_environment( + environment=environment, + identifier="env-1", + resolved_symbol_table_json=symtab_json, + ) + + mock_from_json.assert_called_once_with(symtab_json) + call_kwargs = mock_session_instance.enter_environment.call_args.kwargs + assert call_kwargs["resolved_symtab"] is fake_symtab + + def test_run_task_forwards_resolved_symtab_when_json_present( + self, adapter: RustSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """When resolved_symbol_table_json is provided, it's parsed and forwarded.""" + step_script = MagicMock() + symtab_json = '[{"name":"Job.Name","type":"string","value":"TestJob"}]' + fake_symtab = MagicMock() + + with ( + patch.object(rust_module, "deserialize_step"), + patch.object( + rust_module.SerializedSymbolTable, "from_json_str", return_value=fake_symtab + ) as mock_from_json, + ): + adapter.run_task( + step_script=step_script, + task_parameter_values={}, + resolved_symbol_table_json=symtab_json, + ) + + mock_from_json.assert_called_once_with(symtab_json) + call_kwargs = mock_session_instance.run_task.call_args.kwargs + assert call_kwargs["resolved_symtab"] is fake_symtab + + def test_enter_environment_graceful_degradation_on_malformed_json( + self, adapter: RustSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """Malformed JSON causes a warning log and proceeds with resolved_symtab=None.""" + environment = MagicMock() + + with ( + patch.object(rust_module, "decode_environment_template"), + patch.object(rust_module, "create_environment"), + patch.object( + rust_module.SerializedSymbolTable, + "from_json_str", + side_effect=ValueError("bad json"), + ), + patch.object(rust_module, "logger") as mock_logger, + ): + adapter.enter_environment( + environment=environment, + identifier="env-1", + resolved_symbol_table_json="not valid json", + ) + + mock_logger.warning.assert_called_once() + assert "resolvedSymbolTable" in mock_logger.warning.call_args[0][0] + call_kwargs = mock_session_instance.enter_environment.call_args.kwargs + assert call_kwargs["resolved_symtab"] is None + + def test_exit_environment_forwards_resolved_symtab_when_json_present( + self, adapter: RustSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """When resolved_symbol_table_json is provided, it's parsed and forwarded.""" + symtab_json = '[{"name":"Job.Name","type":"string","value":"TestJob"}]' + fake_symtab = MagicMock() + + with patch.object( + rust_module.SerializedSymbolTable, "from_json_str", return_value=fake_symtab + ) as mock_from_json: + adapter.exit_environment( + identifier="env-1", + resolved_symbol_table_json=symtab_json, + ) + + mock_from_json.assert_called_once_with(symtab_json) + call_kwargs = mock_session_instance.exit_environment.call_args.kwargs + assert call_kwargs["resolved_symtab"] is fake_symtab + + def test_exit_environment_passes_none_when_json_is_none( + self, adapter: RustSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """When resolved_symbol_table_json is None, the parser is not invoked.""" + with patch.object(rust_module.SerializedSymbolTable, "from_json_str") as mock_from_json: + adapter.exit_environment( + identifier="env-1", + resolved_symbol_table_json=None, + ) + + mock_from_json.assert_not_called() + call_kwargs = mock_session_instance.exit_environment.call_args.kwargs + assert call_kwargs["resolved_symtab"] is None + + def test_exit_environment_graceful_degradation_on_malformed_json( + self, adapter: RustSessionRuntime, mock_session_instance: MagicMock + ) -> None: + """Malformed JSON causes a warning log and proceeds with resolved_symtab=None.""" + with ( + patch.object( + rust_module.SerializedSymbolTable, + "from_json_str", + side_effect=ValueError("bad json"), + ), + patch.object(rust_module, "logger") as mock_logger, + ): + adapter.exit_environment( + identifier="env-1", + resolved_symbol_table_json="{not json", + ) + + mock_logger.warning.assert_called_once() + assert "resolvedSymbolTable" in mock_logger.warning.call_args[0][0] + call_kwargs = mock_session_instance.exit_environment.call_args.kwargs + assert call_kwargs["resolved_symtab"] is None diff --git a/test/unit/sessions/test_session.py b/test/unit/sessions/test_session.py index 4b97375d..266fa7e2 100644 --- a/test/unit/sessions/test_session.py +++ b/test/unit/sessions/test_session.py @@ -2663,3 +2663,182 @@ def test_ordinary_exception_does_not_emit_failure_telemetry( crash_session._start_action() mock_telemetry.assert_not_called() + + +class TestResolvedSymbolTableForwarding: + """Tests for resolved_symbol_table_json forwarding through the session layer.""" + + def test_exit_environment_forwards_supplied_table_to_runtime( + self, + session: Session, + mock_runtime: MagicMock, + ) -> None: + """Session.exit_environment forwards a supplied table to the runtime.""" + table = '[{"name":"Job.Name","type":"string","value":"Outer"}]' + # Set up an active environment to exit + from deadline_worker_agent.sessions.session import ActiveEnvironment + + session._active_envs.append( + ActiveEnvironment( + job_env_id="env-1", + session_env_id="session-env-1", + resolved_symbol_table_json='[{"name":"Job.Name","type":"string","value":"Stale"}]', + ) + ) + + session.exit_environment( + job_env_id="env-1", + resolved_symbol_table_json=table, + ) + + mock_runtime.exit_environment.assert_called_once_with( + identifier="session-env-1", + os_env_vars=None, + resolved_symbol_table_json=table, + ) + + def test_exit_environment_falls_back_to_enter_time_table( + self, + session: Session, + mock_runtime: MagicMock, + ) -> None: + """Session.exit_environment uses the stored enter-time table when supplied value is None.""" + enter_table = '[{"name":"Job.Name","type":"string","value":"Outer"}]' + from deadline_worker_agent.sessions.session import ActiveEnvironment + + session._active_envs.append( + ActiveEnvironment( + job_env_id="env-1", + session_env_id="session-env-1", + resolved_symbol_table_json=enter_table, + ) + ) + + session.exit_environment( + job_env_id="env-1", + resolved_symbol_table_json=None, + ) + + mock_runtime.exit_environment.assert_called_once_with( + identifier="session-env-1", + os_env_vars=None, + resolved_symbol_table_json=enter_table, + ) + + def test_exit_environment_supplied_table_takes_precedence( + self, + session: Session, + mock_runtime: MagicMock, + ) -> None: + """A supplied table takes precedence over the stored enter-time table.""" + enter_table = '[{"name":"Job.Name","type":"string","value":"EnterTime"}]' + exit_table = '[{"name":"Job.Name","type":"string","value":"ExitTime"}]' + from deadline_worker_agent.sessions.session import ActiveEnvironment + + session._active_envs.append( + ActiveEnvironment( + job_env_id="env-1", + session_env_id="session-env-1", + resolved_symbol_table_json=enter_table, + ) + ) + + session.exit_environment( + job_env_id="env-1", + resolved_symbol_table_json=exit_table, + ) + + mock_runtime.exit_environment.assert_called_once_with( + identifier="session-env-1", + os_env_vars=None, + resolved_symbol_table_json=exit_table, + ) + + def test_enter_environment_stores_table_on_active_environment( + self, + session: Session, + mock_runtime: MagicMock, + ) -> None: + """Successful enter_environment stores the table on the appended ActiveEnvironment.""" + table = '[{"name":"Job.Name","type":"string","value":"Outer"}]' + mock_runtime.enter_environment.return_value = "session-env-1" + + session.enter_environment( + job_env_id="env-1", + environment=MagicMock(), + resolved_symbol_table_json=table, + ) + + assert len(session._active_envs) == 1 + assert session._active_envs[0].resolved_symbol_table_json == table + + def test_cleanup_forwards_stored_tables_in_reverse_order( + self, + session: Session, + mock_runtime: MagicMock, + ) -> None: + """_cleanup forwards each environment's own stored table and preserves reverse order.""" + outer_table = '[{"name":"Job.Name","type":"string","value":"Outer"}]' + inner_table = '[{"name":"Job.Name","type":"string","value":"Inner"}]' + from deadline_worker_agent.sessions.session import ActiveEnvironment + + session._active_envs = [ + ActiveEnvironment( + job_env_id="env-outer", + session_env_id="session-env-outer", + resolved_symbol_table_json=outer_table, + ), + ActiveEnvironment( + job_env_id="env-inner", + session_env_id="session-env-inner", + resolved_symbol_table_json=inner_table, + ), + ] + + # Mock _monitor_action so cleanup can run without blocking + with patch.object(session, "_monitor_action", return_value=[]): + session._cleanup() + + # exit_environment should be called inner-most first, then outer + calls = mock_runtime.exit_environment.call_args_list + assert len(calls) == 2 + # First call: inner environment with inner table + assert calls[0].kwargs["identifier"] == "session-env-inner" + assert calls[0].kwargs["resolved_symbol_table_json"] == inner_table + # Second call: outer environment with outer table + assert calls[1].kwargs["identifier"] == "session-env-outer" + assert calls[1].kwargs["resolved_symbol_table_json"] == outer_table + + def test_init_passes_resolved_symbol_table_json_to_runtime_config( + self, + asset_sync: MagicMock, + job_details: "JobDetails", + os_user: "SessionUser | None", + mock_create_runtime: MagicMock, + queue_id: str, + session_action_queue: MagicMock, + action_update_callback: MagicMock, + action_update_lock: MagicMock, + session_root_dir: Path, + ) -> None: + """Session.__init__ passes the resolved_symbol_table_json into SessionRuntimeConfig.""" + table = '[{"name":"Job.Name","type":"string","value":"SessionLevel"}]' + + Session( + id="session-symtab-test", + asset_sync=asset_sync, + env=None, + job_details=job_details, + os_user=os_user, + queue=session_action_queue, + queue_id=queue_id, + job_id="job-1234", + action_update_callback=action_update_callback, + action_update_lock=action_update_lock, + session_root_dir=session_root_dir, + resolved_symbol_table_json=table, + ) + + mock_create_runtime.assert_called_once() + config = mock_create_runtime.call_args[0][1] + assert config.resolved_symbol_table_json == table