Pass persisted backend task IDs to Pulsar clients

TES submission returns a provider-assigned task id that Galaxy stores as the job runner external id. Status polling rebuilds the Pulsar client with the Galaxy job id because that id must continue to drive the files and OIDC token endpoints, so the TES client has had no way to recover the provider task id and polls the wrong task.

Pass the recorded external id beside the Galaxy job id when constructing Pulsar clients for polling and cancellation. Older pulsar-galaxy-lib releases safely ignore the extra destination parameter; a paired Pulsar change consumes it for TES get_task and cancel_task operations.

Read the value from the persisted job rather than job_state.job_id, whose fallback to the Galaxy job id would invent a provider id when none was recorded.
This commit is contained in:
John Chilton
2026-08-20 16:06:19 -04:00
parent 65c86d4d92
commit abe22ca628
2 changed files with 126 additions and 6 deletions
+22 -5
View File
@@ -699,10 +699,20 @@ class PulsarJobRunner(AsynchronousJobRunner[AsynchronousJobState]):
def get_client_from_state(self, job_state: AsynchronousJobState) -> "BaseJobClient":
job_destination_params = job_state.job_destination.params
job_id = job_state.job_wrapper.job_id # we want the Galaxy ID here, job_state.job_id is the external one.
return self.get_client(job_destination_params, job_id)
job_wrapper = job_state.job_wrapper
job_id = job_wrapper.job_id # we want the Galaxy ID here, job_state.job_id is the external one.
# Read the external id from the job rather than from job_state.job_id, which falls
# back to the Galaxy id when nothing was ever recorded.
external_id = job_wrapper.get_job().get_job_runner_external_id()
return self.get_client(job_destination_params, job_id, external_id=external_id)
def get_client(self, job_destination_params: dict[str, Any], job_id, env: list | None = None) -> "BaseJobClient":
def get_client(
self,
job_destination_params: dict[str, Any],
job_id,
env: list | None = None,
external_id: str | None = None,
) -> "BaseJobClient":
# Cannot use url_for outside of web thread.
# files_endpoint = url_for( controller="job_files", job_id=encoded_job_id )
if env is None:
@@ -719,6 +729,11 @@ class PulsarJobRunner(AsynchronousJobRunner[AsynchronousJobState]):
get_client_kwds = dict(
job_id=str(job_id), files_endpoint=files_endpoint, token_endpoint=token_endpoint, env=env
)
if external_id:
# TES assigns a task id at submission that cannot be derived from the
# Galaxy job id. Pass it beside the Galaxy id so provider operations use
# the recorded task while Galaxy file and token endpoints remain valid.
get_client_kwds["external_id"] = str(external_id)
# Turn MutableDict into standard dict for pulsar consumption
job_destination_params = dict(job_destination_params.items())
return self.client_manager.get_client(job_destination_params, **get_client_kwds)
@@ -821,7 +836,9 @@ class PulsarJobRunner(AsynchronousJobRunner[AsynchronousJobState]):
if not job.job_runner_external_id:
return
# if our local job has JobExternalOutputMetadata associated, then our primary job has to have already finished
client = self.get_client(job.destination_params, job.job_runner_external_id)
client = self.get_client(
job.destination_params, job.job_runner_external_id, external_id=job.job_runner_external_id
)
job_ext_output_metadata = job.get_external_output_metadata()
if not PulsarJobRunner.__remote_metadata(client) and job_ext_output_metadata:
pid = job_ext_output_metadata[
@@ -858,7 +875,7 @@ class PulsarJobRunner(AsynchronousJobRunner[AsynchronousJobState]):
pulsar_url = job.job_runner_name
job_id = job.job_runner_external_id
log.debug(f"Attempt remote Pulsar kill of job with url {pulsar_url} and id {job_id}")
client = self.get_client(job.destination_params, job_id)
client = self.get_client(job.destination_params, job_id, external_id=job_id)
client.kill()
def recover(self, job: model.Job, job_wrapper: "MinimalJobWrapper") -> None:
+104 -1
View File
@@ -1,6 +1,10 @@
"""Unit tests for Pulsar job runner utility methods."""
"""Unit tests for Pulsar job runner utility methods and client construction."""
from types import SimpleNamespace
from typing import (
Any,
cast,
)
from galaxy.jobs.runners.pulsar import PulsarJobRunner
@@ -58,3 +62,102 @@ def test_rewrite_container_noop_without_container():
# Should not raise when there is no resolved container.
compute_environment = _ComputeEnvironment({IMAGE: REWRITTEN})
PulsarJobRunner._rewrite_container_for_compute_environment(None, compute_environment)
class RecordingClient:
def __init__(self, destination_params, **kwargs):
self.destination_params = destination_params
self.killed = False
for key, value in kwargs.items():
setattr(self, key, value)
def kill(self):
self.killed = True
class RecordingClientManager:
def __init__(self):
self.calls = []
self.clients = []
def get_client(self, destination_params, **kwargs):
self.calls.append((destination_params, kwargs))
client = RecordingClient(destination_params, **kwargs)
self.clients.append(client)
return client
def _runner():
"""A runner with just enough wired up to build clients."""
runner = cast(Any, object.__new__(PulsarJobRunner))
runner.app = SimpleNamespace(
security=SimpleNamespace(encode_id=lambda job_id, kind=None: f"enc{job_id}"),
config=SimpleNamespace(nginx_upload_job_files_path=None),
)
runner.galaxy_url = "http://galaxy.example"
runner.client_manager = RecordingClientManager()
return runner
def _job_state(galaxy_job_id, external_id):
job = SimpleNamespace(get_job_runner_external_id=lambda: external_id)
job_wrapper = SimpleNamespace(job_id=galaxy_job_id, get_job=lambda: job)
return SimpleNamespace(
job_destination=SimpleNamespace(params={"url": "http://pulsar.example"}),
job_wrapper=job_wrapper,
job_id=external_id or str(galaxy_job_id),
)
def test_get_client_omits_external_id_when_absent():
runner = _runner()
runner.get_client({}, 543)
_destination_params, kwargs = runner.client_manager.calls[0]
assert "external_id" not in kwargs
assert kwargs["job_id"] == "543"
def test_get_client_passes_external_id_through():
runner = _runner()
runner.get_client({}, 543, external_id="tes-task-abc")
_destination_params, kwargs = runner.client_manager.calls[0]
assert kwargs["external_id"] == "tes-task-abc"
# The Galaxy id still drives the job files and token endpoints.
assert kwargs["job_id"] == "543"
assert "enc543" in kwargs["files_endpoint"]
def test_get_client_from_state_supplies_the_recorded_external_id():
"""TES status polling has to use the id returned by create_task."""
runner = _runner()
runner.get_client_from_state(_job_state(543, "tes-task-abc"))
_destination_params, kwargs = runner.client_manager.calls[0]
assert kwargs["job_id"] == "543"
assert kwargs["external_id"] == "tes-task-abc"
def test_get_client_from_state_does_not_invent_an_external_id():
"""job_state.job_id falls back to the Galaxy id; that is not a backend name."""
runner = _runner()
runner.get_client_from_state(_job_state(543, None))
_destination_params, kwargs = runner.client_manager.calls[0]
assert "external_id" not in kwargs
def test_stop_job_supplies_recorded_external_id_to_kill_client():
runner = _runner()
external_id = "tes-task-abc"
job = SimpleNamespace(
id=543,
job_runner_external_id=external_id,
job_runner_name="pulsar",
destination_params={"url": "http://pulsar.example", "remote_metadata": True},
get_external_output_metadata=lambda: [],
)
job_wrapper = SimpleNamespace(get_job=lambda: job)
runner.stop_job(job_wrapper)
_destination_params, kill_kwargs = runner.client_manager.calls[-1]
assert kill_kwargs["external_id"] == external_id
assert runner.client_manager.clients[-1].killed