mirror of
https://github.com/galaxyproject/galaxy.git
synced 2026-08-28 18:08:19 +08:00
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:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user