diff --git a/CHANGELOG.md b/CHANGELOG.md index f936c15..92239db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `tilebox-workflows`: Add `workflow` selection by reference or object to job submission and automation creation, + and `workflows` filters to job queries. Omitting scope preserves unscoped direct-runner behavior. +- `tilebox-workflows`: Preserve task workflow IDs and inherit the executing parent's scope for children. + Keep child scope inheritance internal to the runner and task discovery and release content unbound. +- `tilebox-workflows`: Expose workflow IDs, references, project slugs, and presence-aware legacy alias policies. + Retain `.slug` compatibility. + ### Changed +- `tilebox-workflows`: Always create scoped-only workflows with legacy task aliases disabled. +- `tilebox-workflows`: Use `UnboundTaskIdentifier` for task metadata, runner registries, and release contents; + keep execution scope in `TaskIdentifier` and resolve each future task's scope before merging submissions. +- `tilebox-workflows`: Rename the `Workflow` constructor field `slug` to `reference`; retain `.slug` as a read-only + alias. Order workflows by reference, name, and ID (missing IDs first), ignoring other metadata when sorting. - `tilebox-core`: New package to own shared Connect RPC infrastructure and protobuf-py messages, public `tilebox.query` interval/pagination types, and `tilebox.errors` exceptions with a common `TileboxError` base. - `tilebox-grpc`: Retain an empty compatibility package depending on `tilebox-core>=0.67.0` for one release. diff --git a/tilebox-workflows/README.md b/tilebox-workflows/README.md index 1db0107..abbec38 100644 --- a/tilebox-workflows/README.md +++ b/tilebox-workflows/README.md @@ -75,6 +75,41 @@ runner = client.runner(tasks=[MyFirstTask]) runner.run_all() ``` +## Workflow-scoped tasks + +To submit tasks to a managed workflow, pass its remote reference or a `Workflow` object: + +```python +jobs.submit("scoped-job", MyFirstTask(), workflow="my-workflow") + +workflow = client.workflows().find("my-workflow") +jobs.submit("another-job", MyFirstTask(), workflow=workflow) +jobs.query(workflows=[workflow, "tilebox:shared-workflow"]) +``` + +The workflow's immutable UUID scopes each root task. A failed reference lookup raises an error without submitting. +Omitting `workflow` keeps roots unscoped, including jobs submitted from inside another task. Existing direct runners +continue to register and execute unscoped tasks. + +Task metadata, runner registries, and release contents use `UnboundTaskIdentifier` (name and version). +Submissions and executions use `TaskIdentifier`, which also carries the optional workflow ID. Use `.bind(workflow_id)` +to create an execution identifier and `.unbind()` to recover its implementation identifier. + +Children inherit the executing parent's scope. Execution through a legacy unscoped alias therefore produces unscoped +children. The runner handles inheritance internally; `context.submit_subtask()` and `context.submit_subtasks()` +do not expose workflow selection. Cross-workflow and explicitly unscoped child overrides are deferred until the CLI +can distinguish scope-aware workers from workers that need its compatibility adapter. + +Both `create_cron_automation()` and `create_storage_event_automation()` accept the same `workflow` parameter as +`jobs.submit()`. Workflow objects expose `id`, `reference`, `project_slug`, and `legacy_task_aliases`; `.slug` remains +a read-only alias for the remote reference. When constructing a `Workflow` directly, use `reference=` instead of +`slug=`. The project slug is not a unique remote selector. Workflows sort by reference, name, and ID, with missing +IDs first; other metadata does not affect ordering. + +`client.workflows().create(name, summary)` always creates a scoped-only workflow; creating legacy workflows is not +supported. Existing workflows still expose their alias policy. A `None` policy in a response means the server did not +report that field, rather than reporting `False`. + ## Concurrent worker execution A worker runtime can execute multiple tasks concurrently in one Python process. Each execution receives a newly diff --git a/tilebox-workflows/tests/automations/test_client.py b/tilebox-workflows/tests/automations/test_client.py index 772ae14..4304d9c 100644 --- a/tilebox-workflows/tests/automations/test_client.py +++ b/tilebox-workflows/tests/automations/test_client.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock from uuid import UUID, uuid4 +import pytest from hypothesis.stateful import Bundle, RuleBasedStateMachine, consumes, rule from hypothesis.strategies import lists from protobuf.wkt.google.protobuf.empty_pb import Empty @@ -15,7 +16,10 @@ AutomationPrototype, CronTrigger, StorageEventTrigger, + StorageLocation, + StorageType, TaskIdentifier, + Workflow, uuid_to_uuid_message, ) from tilebox.workflows.v1.automation_connect import AutomationServiceClientSync @@ -23,6 +27,9 @@ from tilebox.workflows.v1.automation_pb import Automations, DeleteAutomationRequest from tilebox.workflows.v1.automation_pb import CronTrigger as CronTriggerMessage from tilebox.workflows.v1.automation_pb import StorageEventTrigger as StorageEventTriggerMessage +from tilebox.workflows.v1.workflows_connect import WorkflowsServiceClientSync +from tilebox.workflows.v1.workflows_pb import GetWorkflowRequest +from tilebox.workflows.v1.workflows_pb import Workflow as WorkflowMessage class MockAutomationService: @@ -77,6 +84,44 @@ def list_automations(self, req: Empty) -> Automations: return Automations(automations=list(self.automations.values())) +@pytest.mark.parametrize("kind", ["cron", "storage"]) +@pytest.mark.parametrize("selection", ["reference", "object", "unscoped"]) +def test_automation_scope_survives_creation_and_reads(kind: str, selection: str) -> None: + class ScheduledTask(CronTask): + pass + + class StorageTask(StorageEventTask): + pass + + workflow = Workflow("owner:workflow-123", "Workflow", "", id=uuid4(), project_slug="workflow") + + def find(request: GetWorkflowRequest) -> WorkflowMessage: + if request.workflow_reference != workflow.reference: + raise NotFoundError("Workflow not found") + return workflow.to_message() + + workflow_service = MagicMock(spec=WorkflowsServiceClientSync) + workflow_service.get_workflow.side_effect = find + mock_service = MockAutomationService() + client = AutomationClient(MagicMock(spec=AutomationServiceClientSync, wraps=mock_service), workflow_service) + selected = workflow.reference if selection == "reference" else workflow if selection == "object" else None + if kind == "cron": + created = client.create_cron_automation("Scheduled", ScheduledTask(), "0 * * * *", workflow=selected) + else: + storage = StorageLocation(uuid4(), "test-bucket", StorageType.AWS_S3) + created = client.create_storage_event_automation( + "Storage", StorageTask(), (storage, "*.tif"), workflow=selected + ) + expected = workflow.id if selected is not None else None + assert created.prototype.identifier.workflow_id == expected + assert client.find(created.id).prototype.identifier.workflow_id == expected + assert client.all()[0].prototype.identifier.workflow_id == expected + stored = mock_service.automations[created.id].prototype + assert stored is not None + assert stored.identifier is not None + assert stored.identifier.workflow_id == uuid_to_uuid_message(expected) + + class AutomationCRUDOperations(RuleBasedStateMachine): """ A state machine that tests the CRUD operations of the Automations client. diff --git a/tilebox-workflows/tests/jobs/test_client.py b/tilebox-workflows/tests/jobs/test_client.py index eb12773..14cb117 100644 --- a/tilebox-workflows/tests/jobs/test_client.py +++ b/tilebox-workflows/tests/jobs/test_client.py @@ -18,6 +18,7 @@ LogSeverity, LogSeverityLiteral, Span, + Workflow, uuid_message_to_uuid, uuid_to_uuid_message, ) @@ -40,6 +41,9 @@ VisualizeJobRequest, ) from tilebox.workflows.v1.telemetry_pb import PaginatedLogsData, PaginatedSpansData +from tilebox.workflows.v1.workflows_connect import WorkflowsServiceClientSync +from tilebox.workflows.v1.workflows_pb import GetWorkflowRequest +from tilebox.workflows.v1.workflows_pb import Workflow as WorkflowMessage class DummyTask(Task): @@ -190,8 +194,10 @@ class MockJobService: def __init__(self) -> None: self.jobs: dict[UUID, JobMessage] = {} self.query_requests: list[QueryJobsRequest] = [] + self.submit_requests: list[SubmitJobRequest] = [] def submit_job(self, req: SubmitJobRequest) -> JobMessage: + self.submit_requests.append(req) job_id = uuid4() job = JobMessage( id=uuid_to_uuid_message(job_id), @@ -246,16 +252,83 @@ def query_jobs(self, req: QueryJobsRequest) -> QueryJobsResponse: return QueryJobsResponse(jobs=list(self.jobs.values())) -def test_query_filters_by_clusters() -> None: +@pytest.mark.parametrize("selection", ["reference", "object", "unscoped"]) +def test_submit_workflow_scope(selection: str) -> None: + class AnotherTask(Task): + def execute(self, context: ExecutionContext) -> None: + _ = context + + workflow = Workflow("owner:workflow-123", "Workflow", "", id=uuid4(), project_slug="workflow") + + def find(request: GetWorkflowRequest) -> WorkflowMessage: + if request.workflow_reference != workflow.reference: + raise NotFoundError("Workflow not found") + return workflow.to_message() + + workflow_service = MagicMock(spec=WorkflowsServiceClientSync) + workflow_service.get_workflow.side_effect = find + mock_service = MockJobService() + client = JobClient( + MagicMock(spec=JobServiceClientSync, wraps=mock_service), MagicMock(), NoopWorkflowTracer(), workflow_service + ) + selected = workflow.reference if selection == "reference" else workflow if selection == "object" else None + client.submit("scoped job", [DummyTask(), AnotherTask()], cluster=["a", "b"], workflow=selected) + submissions = mock_service.submit_requests[-1].tasks + assert submissions is not None + assert len(submissions.identifier_lookup) == 2 + assert {uuid_message_to_optional_uuid(identifier.workflow_id) for identifier in submissions.identifier_lookup} == { + workflow.id if selected is not None else None + } + # A subsequent ordinary submission must not inherit a previous submission's scope. + client.submit("direct job", DummyTask()) + direct = mock_service.submit_requests[-1].tasks + assert direct is not None + assert direct.identifier_lookup[0].workflow_id is None + + +@pytest.mark.parametrize("selection", ["missing", "empty", "missing-id", "nil-id"]) +def test_submit_invalid_workflow_never_falls_back_to_unscoped(selection: str) -> None: + mock_service = MockJobService() + workflow_service = MagicMock(spec=WorkflowsServiceClientSync) + workflow_service.get_workflow.side_effect = NotFoundError("Workflow not found") + selected = { + "missing": "missing", + "empty": "", + "missing-id": Workflow("old", "Old", ""), + "nil-id": Workflow("invalid", "Invalid", "", id=UUID(int=0)), + }[selection] + client = JobClient( + MagicMock(spec=JobServiceClientSync, wraps=mock_service), MagicMock(), NoopWorkflowTracer(), workflow_service + ) + with pytest.raises((NotFoundError, ValueError)): + client.submit("must not submit", DummyTask(), workflow=selected) + assert not mock_service.jobs + assert not mock_service.submit_requests + + +@pytest.mark.parametrize( + ("workflows", "expected"), + [ + (None, []), + ([], []), + ("owner:one", ["owner:one"]), + (Workflow("one-123", "One", "", project_slug="one"), ["one-123"]), + (["owner:one", Workflow("two-456", "Two", "", project_slug="two")], ["owner:one", "two-456"]), + ], +) +def test_query_filters_by_clusters_and_workflows( + workflows: Workflow | str | list[Workflow | str] | None, expected: list[str] +) -> None: mock_service = MockJobService() service = MagicMock(spec=JobServiceClientSync, wraps=mock_service) job_client = JobClient(service, MagicMock(), NoopWorkflowTracer()) - job_client.query((uuid4(), uuid4()), clusters=["cluster-a", "cluster-b"]) + job_client.query((uuid4(), uuid4()), clusters=["cluster-a", "cluster-b"], workflows=workflows) filters = mock_service.query_requests[-1].filters assert filters is not None assert list(filters.cluster_slugs) == ["cluster-a", "cluster-b"] + assert list(filters.workflow_references) == expected @pytest.mark.parametrize("use_interval", [False, True]) diff --git a/tilebox-workflows/tests/runner/test_worker_concurrency.py b/tilebox-workflows/tests/runner/test_worker_concurrency.py index 1769a25..b2e74e1 100644 --- a/tilebox-workflows/tests/runner/test_worker_concurrency.py +++ b/tilebox-workflows/tests/runner/test_worker_concurrency.py @@ -11,6 +11,7 @@ import pytest from protobuf.wkt.google.protobuf.empty_pb import Empty +from _tilebox.uuid import uuid_message_to_optional_uuid from tilebox.datasets.uuid import must_uuid_to_uuid_message from tilebox.workflows import ExecutionContext, Runner, Task from tilebox.workflows.cache import InMemoryCache, JobCache @@ -23,6 +24,7 @@ from tilebox.workflows.task import TaskMeta from tilebox.workflows.v1 import core_pb, worker_pb from tilebox.workflows.v1.worker_connect import WorkerServiceClientSync +from tilebox.workflows.v1.workflows_pb import Workflow as WorkflowMessage def test_worker_executes_tasks_concurrently_with_isolated_execution_state( @@ -127,6 +129,108 @@ async def execute(self, context: ExecutionContext) -> None: assert {attributes["task_id"] for attributes in log_attributes} == task_ids +def test_worker_preserves_execution_scope_across_concurrent_descendants( # noqa: PLR0915 + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + scopes = [uuid4(), uuid4(), None] + barrier = threading.Barrier(len(scopes)) + + class Grandchild(Task): + def execute(self, context: ExecutionContext) -> None: + context.current_task.display = str(context.current_task.identifier.workflow_id) + + class Child(Task): + def execute(self, context: ExecutionContext) -> None: + future = context.submit_subtask(Grandchild()) + assert future.identifier().workflow_id == context.current_task.identifier.workflow_id + + class Parent(Task): + def execute(self, context: ExecutionContext) -> None: + barrier.wait(timeout=5) + context.current_task.display = str(context.current_task.identifier.workflow_id) + futures = [context.submit_subtask(Child()), *context.submit_subtasks([Child(), Child()])] + assert all( + future.identifier().workflow_id == context.current_task.identifier.workflow_id for future in futures + ) + + fake_client = MagicMock() + fake_client._tracer = NoopWorkflowTracer() + fake_client._auth = {"url": "https://worker.example", "token": "worker-key"} + fake_client._task_logger = StructuredLogger(logging.getLogger("tilebox.workflows.tests.scoped-worker")) + monkeypatch.setenv("TILEBOX_RUNTIME_DIR", str(tmp_path)) + monkeypatch.setenv("TILEBOX_RUNTIME_TOKEN", "test-token") + server = WorkerServer(Runner(tasks=[Parent, Child, Grandchild]), "127.0.0.1:0") + headers = {"authorization": "Bearer test-token"} + + with ( + patch("tilebox.workflows.runner.worker_service.Client", return_value=fake_client), + patch("tilebox.workflows.runner.worker_service.initialize_logging"), + ): + server.start() + try: + with WorkerServiceClientSync(f"http://127.0.0.1:{server._port}", timeout_ms=10000) as worker: + discovered = worker.list_registered_tasks(Empty(), headers=headers) + assert len(discovered.identifiers) == 3 + assert all( + isinstance(identifier, core_pb.UnboundTaskIdentifier) for identifier in discovered.identifiers + ) + worker.initialize_worker( + worker_pb.InitializeRunnerRequest( + runner_id=must_uuid_to_uuid_message(uuid4()), + # Release metadata must never override the executing task's scope, especially legacy aliases. + workflow=WorkflowMessage(id=must_uuid_to_uuid_message(uuid4()), legacy_task_aliases=True), + ), + headers=headers, + ) + job = _job() + parents = [_task_message(Parent(), job) for _ in scopes] + for parent, scope in zip(parents, scopes, strict=True): + assert parent.identifier is not None + parent.identifier.workflow_id = must_uuid_to_uuid_message(scope) if scope is not None else None + with ThreadPoolExecutor(max_workers=len(scopes)) as pool: + futures = [pool.submit(worker.execute_task, parent, headers=headers) for parent in parents] + responses = [future.result(timeout=10) for future in futures] + for response, scope in zip(responses, scopes, strict=True): + computed = response.computed_task + assert computed is not None + assert computed.display == str(scope) + children = computed.sub_tasks + assert children is not None + child_scopes = [] + for group in children.task_groups: + for task_input, pointer in zip(group.inputs, group.identifier_pointers, strict=True): + identifier = children.identifier_lookup[pointer] + child_scope = uuid_message_to_optional_uuid(identifier.workflow_id) + child_scopes.append(child_scope) + child = core_pb.Task( + id=must_uuid_to_uuid_message(uuid4()), + identifier=identifier, + input=task_input, + job=job.to_message(), + ) + result = worker.execute_task(child, headers=headers).computed_task + assert result is not None + assert result.sub_tasks is not None + grandchildren = result.sub_tasks + grandchild_identifier = grandchildren.identifier_lookup[0] + assert uuid_message_to_optional_uuid(grandchild_identifier.workflow_id) == child_scope + grandchild = core_pb.Task( + id=must_uuid_to_uuid_message(uuid4()), + identifier=grandchild_identifier, + input=grandchildren.task_groups[0].inputs[0], + job=job.to_message(), + ) + final = worker.execute_task(grandchild, headers=headers).computed_task + assert final is not None + assert final.display == str(child_scope) + assert child_scopes == [scope, scope, scope] + assert worker.list_registered_tasks(Empty(), headers=headers) == discovered + worker.shutdown_worker(Empty(), headers=headers) + finally: + server.shutdown() + server.wait() + + def test_lazy_storage_locations_are_loaded_once_during_concurrent_access() -> None: storage_location = MagicMock() storage_location.id = UUID(int=1) @@ -180,7 +284,7 @@ def _job() -> Job: def _task_message(task: Task, job: Job) -> core_pb.Task: - identifier = TaskMeta.for_task(task).identifier + identifier = TaskMeta.for_task(task).identifier.bind() return TaskData( id=uuid4(), identifier=identifier, diff --git a/tilebox-workflows/tests/tasks_data.py b/tilebox-workflows/tests/tasks_data.py index ba71ccf..29a2482 100644 --- a/tilebox-workflows/tests/tasks_data.py +++ b/tilebox-workflows/tests/tasks_data.py @@ -102,7 +102,7 @@ def release_contents(draw: DrawFn) -> ReleaseContent: fingerprint = draw(text(alphabet="abcdef0123456789", min_size=64, max_size=64)) return ReleaseContent( fingerprint=fingerprint, - tasks=draw(lists(task_identifiers(), min_size=1, max_size=3)), + tasks=[identifier.unbind() for identifier in draw(lists(task_identifiers(), min_size=1, max_size=3))], files=draw(lists(release_filesystem_nodes(), min_size=0, max_size=3)), runner_object_path=draw(alphanumerical_text()), command_override=draw(lists(alphanumerical_text(), min_size=0, max_size=3)), @@ -125,10 +125,13 @@ def workflow_releases(draw: DrawFn) -> WorkflowRelease: def workflows(draw: DrawFn) -> Workflow: """A hypothesis strategy for generating random workflows""" return Workflow( - slug=draw(alphanumerical_text(min_size=4, max_size=20)), + reference=draw(alphanumerical_text(min_size=4, max_size=20)), name=draw(alphanumerical_text()), description=draw(alphanumerical_text() | none()) or "", releases=draw(lists(workflow_releases(), min_size=0, max_size=3)), + id=draw(uuids(version=4) | none()), + project_slug=draw(alphanumerical_text()), + legacy_task_aliases=draw(booleans() | none()), ) @@ -145,7 +148,7 @@ def progress_indicators(draw: DrawFn) -> ProgressIndicator: def tasks(draw: DrawFn) -> Task: """A hypothesis strategy for generating random tasks""" task_id = draw(uuids(version=4)) - identifier = draw(task_identifiers()) + identifier = draw(task_identifiers(scoped=True)) state = draw(sampled_from(TaskState)) task_input = draw(task_inputs()) display = draw(alphanumerical_text()) @@ -170,12 +173,13 @@ def idling_responses(draw: DrawFn) -> Idling: @composite -def task_identifiers(draw: DrawFn) -> TaskIdentifier: +def task_identifiers(draw: DrawFn, *, scoped: bool = False) -> TaskIdentifier: """A hypothesis strategy for generating random task_identifiers""" task_name = draw(alphanumerical_text()) major_version = draw(integers(min_value=0, max_value=100)) minor_version = draw(integers(min_value=0, max_value=100)) - return TaskIdentifier(task_name, f"v{major_version}.{minor_version}") + workflow_id = draw(uuids(version=4) | none()) if scoped else None + return TaskIdentifier(task_name, f"v{major_version}.{minor_version}", workflow_id) @composite @@ -288,7 +292,7 @@ def task_submissions(draw: DrawFn) -> TaskSubmissions: """A hypothesis strategy for generating random task_submissions""" task_groups = draw(lists(task_submission_groups(), min_size=1, max_size=10)) cluster_slug_lookup = draw(lists(alphanumerical_text(), min_size=1, max_size=10)) - identifier_lookup = draw(lists(task_identifiers(), min_size=1, max_size=10)) + identifier_lookup = draw(lists(task_identifiers(scoped=True), min_size=1, max_size=10)) display_lookup = draw(lists(alphanumerical_text(), min_size=1, max_size=10)) return TaskSubmissions(task_groups, cluster_slug_lookup, identifier_lookup, display_lookup) @@ -298,7 +302,7 @@ def task_submissions(draw: DrawFn) -> TaskSubmissions: def single_task_submissions(draw: DrawFn) -> SingleTaskSubmission: """A hypothesis strategy for generating random sub_tasks""" cluster_slug = str(draw(uuids(version=4))) - identifier = draw(task_identifiers()) + identifier = draw(task_identifiers(scoped=True)) inputs = draw(task_inputs()) dependencies: list[int] = draw(lists(integers(min_value=0, max_value=1000), min_size=0, max_size=10)) display = draw(alphanumerical_text()) @@ -431,5 +435,8 @@ def query_filters(draw: DrawFn) -> QueryFilters: task_states = list(set(task_states)) # de-duplicate cluster_slugs = draw(lists(alphanumerical_text(min_size=4, max_size=20), min_size=0, max_size=3, unique=True)) + workflow_references = draw(lists(alphanumerical_text(min_size=4, max_size=20), max_size=3, unique=True)) - return QueryFilters(time_interval, id_interval, automation_ids, job_states, name, task_states, cluster_slugs) + return QueryFilters( + time_interval, id_interval, automation_ids, job_states, name, task_states, cluster_slugs, workflow_references + ) diff --git a/tilebox-workflows/tests/test_data.py b/tilebox-workflows/tests/test_data.py index 6e55257..1e9f211 100644 --- a/tilebox-workflows/tests/test_data.py +++ b/tilebox-workflows/tests/test_data.py @@ -1,11 +1,13 @@ from dataclasses import replace -from itertools import product +from itertools import pairwise, product from uuid import UUID import pytest from hypothesis import given +from hypothesis.strategies import none, uuids from protobuf import Oneof +from _tilebox.proto.tilebox.v1.id_pb import ID from tests.tasks_data import ( artifacts, automations, @@ -55,17 +57,40 @@ TaskLease, TaskSubmissionGroup, TaskSubmissions, + UnboundTaskIdentifier, Workflow, WorkflowPresentation, WorkflowRelease, ) from tilebox.workflows.formatting.job import JobWidget, RichDisplayJob from tilebox.workflows.v1 import workflows_pb - - -@given(task_identifiers()) -def test_task_identifiers_to_message_and_back(task_id: TaskIdentifier) -> None: - assert TaskIdentifier.from_message(task_id.to_message()) == task_id +from tilebox.workflows.v1.core_pb import TaskIdentifier as TaskIdentifierMessage +from tilebox.workflows.v1.core_pb import UnboundTaskIdentifier as UnboundTaskIdentifierMessage + + +@given(task_identifiers(), uuids(version=4) | none()) +def test_task_identifiers_to_message_and_back(task_id: TaskIdentifier, scope: UUID | None) -> None: + scoped = replace(task_id, workflow_id=scope) + message = scoped.to_message() + assert message.workflow_id == (ID(uuid=scope.bytes) if scope is not None else None) + assert TaskIdentifier.from_message(message) == scoped + unbound = UnboundTaskIdentifier(task_id.name, task_id.version) + assert scoped.unbind() == unbound + assert unbound.bind(scope) == scoped + assert unbound.bind() == task_id + assert unbound.to_message() == UnboundTaskIdentifierMessage(name=task_id.name, version=task_id.version) + assert UnboundTaskIdentifier.from_message(unbound.to_message()) == unbound + if scope is not None: + assert len({task_id, scoped}) == 2 + assert sorted([scoped, task_id]) == [task_id, scoped] + + +@pytest.mark.parametrize("raw_id", [b"", bytes(16), b"invalid"]) +def test_task_identifiers_reject_invalid_explicit_scope(raw_id: bytes) -> None: + with pytest.raises(ValueError, match=r"Nil|16-char"): + TaskIdentifier.from_message(TaskIdentifierMessage(name="Task", version="v0.0", workflow_id=ID(uuid=raw_id))) + with pytest.raises(ValueError, match="Nil"): + TaskIdentifier("Task", "v0.0", UUID(int=0)) @given(progress_indicators()) @@ -173,6 +198,40 @@ def test_workflow_releases_to_message_and_back(release: WorkflowRelease) -> None @given(workflows()) def test_workflows_to_message_and_back(workflow: Workflow) -> None: assert Workflow.from_message(workflow.to_message()) == workflow + assert workflow.slug == workflow.reference + assert workflow.to_message().reference == workflow.reference + + +def test_workflow_ordering() -> None: + workflow = Workflow(reference="a", name="A", description="") + ordered = [ + workflow, + replace(workflow, id=UUID(int=1)), + replace(workflow, id=UUID(int=2)), + replace(workflow, name="B"), + replace(workflow, reference="b", name="0"), + ] + unordered = list(reversed(ordered)) + assert sorted(unordered) == ordered + for earlier, later in pairwise(ordered): + assert earlier < later + assert earlier <= later + assert later > earlier + assert later >= earlier + + +@given(workflows(), workflows()) +def test_workflow_ordering_ignores_metadata(first: Workflow, second: Workflow) -> None: + second = replace(second, reference=first.reference, name=first.name, id=first.id) + assert not first < second + assert not first > second + assert first <= second + assert first >= second + assert not second < first + assert not second > first + assert second <= first + assert second >= first + assert first != replace(first, description=first.description + "changed") @pytest.mark.parametrize( diff --git a/tilebox-workflows/tests/test_task.py b/tilebox-workflows/tests/test_task.py index e3a6557..8b187cd 100644 --- a/tilebox-workflows/tests/test_task.py +++ b/tilebox-workflows/tests/test_task.py @@ -2,12 +2,14 @@ from collections.abc import Awaitable from dataclasses import dataclass, field, fields, is_dataclass from typing import Annotated +from uuid import uuid4 import pytest from tests.proto.test_pb import SampleArgs from tilebox.workflows.cache import InMemoryCache -from tilebox.workflows.data import TaskIdentifier +from tilebox.workflows.data import Task as TaskData +from tilebox.workflows.data import TaskIdentifier, UnboundTaskIdentifier from tilebox.workflows.runner.task_runner import ExecutionContext as RunnerExecutionContext from tilebox.workflows.task import ( ExecutionContext, @@ -47,7 +49,7 @@ def test_task_validation_simple_task() -> None: class SimpleTask(Task): pass - expected_identifier = TaskIdentifier("SimpleTask", "v0.0") + expected_identifier = UnboundTaskIdentifier("SimpleTask", "v0.0") assert TaskMeta.for_task(SimpleTask).identifier == expected_identifier assert TaskMeta.for_task(SimpleTask).executable is False @@ -58,7 +60,7 @@ class SimpleTask(Task): def identifier() -> tuple[str, str]: return "tilebox.tests.SimpleTask", "v3.2" - expected_identifier = TaskIdentifier("tilebox.tests.SimpleTask", "v3.2") + expected_identifier = UnboundTaskIdentifier("tilebox.tests.SimpleTask", "v3.2") assert TaskMeta.for_task(SimpleTask).identifier == expected_identifier @@ -184,7 +186,7 @@ class SimpleTaskWithClassmethodIdentifier(Task): def identifier(cls) -> tuple[str, str]: return "tilebox.tests.SimpleTaskWithClassmethodIdentifier", "v3.2" - expected_identifier = TaskIdentifier("tilebox.tests.SimpleTaskWithClassmethodIdentifier", "v3.2") + expected_identifier = UnboundTaskIdentifier("tilebox.tests.SimpleTaskWithClassmethodIdentifier", "v3.2") assert TaskMeta.for_task(SimpleTaskWithClassmethodIdentifier).identifier == expected_identifier @@ -520,8 +522,13 @@ def execute(self, context: ExecutionContext) -> None: pass -def test_merge_future_tasks_to_submissions() -> None: - context = RunnerExecutionContext(None, None, job_cache=InMemoryCache()) # ty: ignore[invalid-argument-type] +@pytest.fixture +def context() -> RunnerExecutionContext: + parent = TaskData(uuid4(), TaskIdentifier("Parent", "v0.0")) + return RunnerExecutionContext(None, parent, job_cache=InMemoryCache()) # ty: ignore[invalid-argument-type] + + +def test_merge_future_tasks_to_submissions(context: RunnerExecutionContext) -> None: tasks_1 = context.submit_subtasks([TaskA(3, "three"), TaskA(4, "four"), TaskA(5, "five")]) tasks_2 = context.submit_subtasks([TaskB(3.2), TaskB(3.44), TaskB(3.55)], max_retries=1) tasks_3 = context.submit_subtasks([TaskA(6, "six"), TaskB(8.12)], cluster="other") @@ -547,12 +554,11 @@ def test_merge_future_tasks_to_submissions() -> None: assert group.max_retries_values == [0, 0, 0, 1, 1, 1, 0, 0] assert submissions.cluster_slug_lookup == ["test", "other"] - assert submissions.identifier_lookup == [TaskMeta.for_task(TaskA).identifier, TaskMeta.for_task(TaskB).identifier] + assert submissions.identifier_lookup == [TaskIdentifier("TaskA", "v0.0"), TaskIdentifier("TaskB", "v0.0")] assert submissions.display_lookup == ["TaskA", "TaskB"] -def test_merge_future_tasks_to_submissions_dependencies() -> None: - context = RunnerExecutionContext(None, None, job_cache=InMemoryCache()) # ty: ignore[invalid-argument-type] +def test_merge_future_tasks_to_submissions_dependencies(context: RunnerExecutionContext) -> None: tasks_1 = context.submit_subtasks([TaskA(2, "two"), TaskA(3, "three")]) tasks_2 = context.submit_subtasks([TaskA(4, "four"), TaskA(5, "five")]) tasks_3 = context.submit_subtasks([TaskB(3.2)], depends_on=tasks_1) @@ -579,8 +585,7 @@ def test_merge_future_tasks_to_submissions_dependencies() -> None: assert submissions.task_groups[4].inputs == [serialize_task(TaskB(3.55))] -def test_merge_future_tasks_to_submissions_many_tasks() -> None: - context = RunnerExecutionContext(None, None, job_cache=InMemoryCache()) # ty: ignore[invalid-argument-type] +def test_merge_future_tasks_to_submissions_many_tasks(context: RunnerExecutionContext) -> None: n = 100 tasks_1 = context.submit_subtasks([TaskA(i, f"Task {i}") for i in range(n)]) tasks_2 = context.submit_subtasks([TaskB(i / 3) for i in range(n)], depends_on=tasks_1) @@ -594,8 +599,9 @@ def test_merge_future_tasks_to_submissions_many_tasks() -> None: assert submissions.task_groups[1].identifier_pointers == [1] * n -def test_merge_future_tasks_to_submissions_many_non_mergeable_dependency_groups() -> None: - context = RunnerExecutionContext(None, None, job_cache=InMemoryCache()) # ty: ignore[invalid-argument-type] +def test_merge_future_tasks_to_submissions_many_non_mergeable_dependency_groups( + context: RunnerExecutionContext, +) -> None: n = 100 for i in range(n): task_1 = context.submit_subtasks([TaskA(i, f"Task {i}")]) @@ -606,8 +612,7 @@ def test_merge_future_tasks_to_submissions_many_non_mergeable_dependency_groups( assert len(submissions.task_groups) == 2 * n -def test_merge_future_tasks_two_separate_branches() -> None: - context = RunnerExecutionContext(None, None, job_cache=InMemoryCache()) # ty: ignore[invalid-argument-type] +def test_merge_future_tasks_two_separate_branches(context: RunnerExecutionContext) -> None: task_a = context.submit_subtasks([TaskA(0, "Task 0")]) # left branch task_b_left = context.submit_subtasks([TaskB(0.0)], depends_on=task_a) diff --git a/tilebox-workflows/tests/workflows/test_client.py b/tilebox-workflows/tests/workflows/test_client.py index d4c23b7..bbd09cb 100644 --- a/tilebox-workflows/tests/workflows/test_client.py +++ b/tilebox-workflows/tests/workflows/test_client.py @@ -10,7 +10,7 @@ Cluster, FilesystemNode, ReleaseContent, - TaskIdentifier, + UnboundTaskIdentifier, Workflow, WorkflowRelease, ) @@ -49,10 +49,20 @@ def __init__(self) -> None: self.deploy_requests: list[DeployWorkflowReleaseRequest] = [] self.undeploy_requests: list[UndeployWorkflowReleaseRequest] = [] self.update_requests: list[UpdateWorkflowRequest] = [] + self.create_requests: list[CreateWorkflowRequest] = [] def create_workflow(self, req: CreateWorkflowRequest) -> WorkflowMessage: + self.create_requests.append(req) slug = req.name.lower().replace(" ", "-") - workflow = WorkflowMessage(reference=slug, name=req.name, description=req.description, summary=req.summary) + workflow = WorkflowMessage( + id=uuid_to_uuid_message(uuid4()), + reference=slug, + name=req.name, + description=req.description, + summary=req.summary, + project_slug=req.project_slug, + legacy_task_aliases=req.legacy_task_aliases if req.has_field("legacy_task_aliases") else True, + ) self.workflows[workflow.reference] = workflow return workflow @@ -105,15 +115,23 @@ def undeploy_workflow_release(self, req: UndeployWorkflowReleaseRequest) -> Unde def test_create_workflow() -> None: - service = MagicMock(spec=WorkflowsServiceClientSync, wraps=MockWorkflowService()) + mock_service = MockWorkflowService() + service = MagicMock(spec=WorkflowsServiceClientSync, wraps=mock_service) workflow_client = WorkflowClient(service) - workflow = workflow_client.create("Agentic Workflow", "Short summary", "Description") + workflow = workflow_client.create("Agentic Workflow", "Short summary", "Description", project_slug="agentic") assert workflow.slug == "agentic-workflow" assert workflow.name == "Agentic Workflow" assert workflow.description == "Description" assert workflow.summary == "Short summary" + assert workflow.id is not None + assert workflow.reference == workflow.slug + assert workflow.project_slug == "agentic" + assert workflow.legacy_task_aliases is False + request = mock_service.create_requests[-1] + assert request.has_field("legacy_task_aliases") + assert request.legacy_task_aliases is False def test_list_find_and_delete_workflows() -> None: @@ -237,7 +255,7 @@ def test_workflow_release_data_includes_release_content_and_clusters() -> None: artifact=Artifact(artifact_id, "a" * 64), content=ReleaseContent( fingerprint="b" * 64, - tasks=[TaskIdentifier("tilebox.com/task/Review", "v1.0")], + tasks=[UnboundTaskIdentifier("tilebox.com/task/Review", "v1.0")], files=[FilesystemNode(".", directory=True, children=[FilesystemNode("main.py")])], runner_object_path="my_module.my_runner:runner", command_override=["python", "main.py"], diff --git a/tilebox-workflows/tilebox/workflows/automations/client.py b/tilebox-workflows/tilebox/workflows/automations/client.py index 73b068b..ef1b497 100644 --- a/tilebox-workflows/tilebox/workflows/automations/client.py +++ b/tilebox-workflows/tilebox/workflows/automations/client.py @@ -16,11 +16,16 @@ from tilebox.workflows.task import _task_meta from tilebox.workflows.v1.automation_connect import AutomationServiceClientSync from tilebox.workflows.v1.automation_pb import DeleteAutomationRequest +from tilebox.workflows.v1.workflows_connect import WorkflowsServiceClientSync +from tilebox.workflows.workflows.client import WorkflowReferenceLike, resolve_workflow_id class AutomationClient: - def __init__(self, service: AutomationServiceClientSync) -> None: + def __init__( + self, service: AutomationServiceClientSync, workflow_service: WorkflowsServiceClientSync | None = None + ) -> None: self._service = service + self._workflow_service = workflow_service def all(self) -> list[AutomationPrototype]: """List all registered automations. @@ -44,13 +49,15 @@ def find(self, automation_id: UUID | str) -> AutomationPrototype: automation_id = UUID(automation_id) return AutomationPrototype.from_message(self._service.get_automation(must_uuid_to_uuid_message(automation_id))) - def create_cron_automation( + def create_cron_automation( # noqa: PLR0913 self, name: str, task: CronTask, cron_schedules: str | list[str], cluster: ClusterSlugLike | None = None, max_retries: int = 0, + *, + workflow: WorkflowReferenceLike | None = None, ) -> AutomationPrototype: """Create a new automation that is triggered by cron schedules. @@ -60,6 +67,7 @@ def create_cron_automation( cron_schedules: The cron schedules to trigger the task. cluster: The cluster to run the task on. If not provided, the default cluster will be used. max_retries: The maximum number of retries for the task in case of failure. Defaults to 0. + workflow: Workflow reference or object to scope the prototype to. Omission leaves it unscoped. Returns: The created automation. @@ -70,12 +78,13 @@ def create_cron_automation( if not cron_schedules: raise ValueError("At least one cron trigger schedule must be provided.") + workflow_id = resolve_workflow_id(workflow, self._workflow_service) automation = AutomationPrototype( id=UUID(int=0), name=name, prototype=SingleTaskSubmission( cluster_slug=to_cluster_slug(cluster or ""), - identifier=_task_meta(task).identifier, + identifier=_task_meta(task).identifier.bind(workflow_id), input=task._serialize_args(), # noqa: SLF001 dependencies=[], display=task.__class__.__name__, @@ -86,13 +95,15 @@ def create_cron_automation( ) return AutomationPrototype.from_message(self._service.create_automation(automation.to_message())) - def create_storage_event_automation( + def create_storage_event_automation( # noqa: PLR0913 self, name: str, task: StorageEventTask, triggers: list[tuple[StorageLocation, str]] | tuple[StorageLocation, str], cluster: ClusterSlugLike | None = None, max_retries: int = 0, + *, + workflow: WorkflowReferenceLike | None = None, ) -> AutomationPrototype: """Create a new automation that is triggered by an object being added to a storage location. @@ -102,6 +113,7 @@ def create_storage_event_automation( triggers: Tuples of storage location and glob pattern to trigger the task. cluster: The cluster to run the task on. If not provided, the default cluster will be used. max_retries: The maximum number of retries for the task in case of failure. Defaults to 0. + workflow: Workflow reference or object to scope the prototype to. Omission leaves it unscoped. Returns: The created automation. @@ -112,12 +124,13 @@ def create_storage_event_automation( if not triggers: raise ValueError("At least one bucket trigger must be provided.") + workflow_id = resolve_workflow_id(workflow, self._workflow_service) automation = AutomationPrototype( id=UUID(int=0), name=name, prototype=SingleTaskSubmission( cluster_slug=to_cluster_slug(cluster or ""), - identifier=_task_meta(task).identifier, + identifier=_task_meta(task).identifier.bind(workflow_id), input=task._serialize_args(), # noqa: SLF001 dependencies=[], display=task.__class__.__name__, diff --git a/tilebox-workflows/tilebox/workflows/client.py b/tilebox-workflows/tilebox/workflows/client.py index 5a47ed2..3a5cd81 100644 --- a/tilebox-workflows/tilebox/workflows/client.py +++ b/tilebox-workflows/tilebox/workflows/client.py @@ -115,7 +115,7 @@ def jobs(self) -> JobClient: A client for the jobs service. """ - return JobClient(self._job_service, self._telemetry_service, self._tracer) + return JobClient(self._job_service, self._telemetry_service, self._tracer, self._workflow_service) def runner( self, @@ -188,7 +188,7 @@ def automations(self) -> AutomationClient: Returns: A client for the automations service. """ - return AutomationClient(self._automation_service) + return AutomationClient(self._automation_service, self._workflow_service) def storage_locations(self) -> StorageLocationClient: """Get a client for storage locations.""" diff --git a/tilebox-workflows/tilebox/workflows/data.py b/tilebox-workflows/tilebox/workflows/data.py index 3472295..ac217d9 100644 --- a/tilebox-workflows/tilebox/workflows/data.py +++ b/tilebox-workflows/tilebox/workflows/data.py @@ -55,17 +55,71 @@ class TaskState(Enum): _TASK_STATES = {state.value: state for state in TaskState} -@dataclass(order=True, eq=True, unsafe_hash=True, frozen=True) -class TaskIdentifier: +@dataclass(order=True, frozen=True) +class UnboundTaskIdentifier: + """A task implementation's name and version, independent of workflow scope.""" + name: str version: str @classmethod - def from_message(cls, identifier: core_pb.TaskIdentifier | core_pb.UnboundTaskIdentifier) -> Self: + def from_message(cls, identifier: core_pb.UnboundTaskIdentifier) -> Self: return cls(name=identifier.name, version=identifier.version) + def to_message(self) -> core_pb.UnboundTaskIdentifier: + return core_pb.UnboundTaskIdentifier(name=self.name, version=self.version) + + def bind(self, workflow_id: UUID | None = None) -> "TaskIdentifier": + """Create an execution identifier; None selects unscoped execution.""" + return TaskIdentifier(self.name, self.version, workflow_id) + + @classmethod + def from_name_and_version(cls, name: str, version: str) -> Self: + """Construct an implementation identifier, validating its name and version.""" + if name == "": + raise ValueError("A task name is required") + if len(name) > 256: + raise ValueError("The task name is too long") + _parse_version(version) + return cls(name=name, version=version) + + +@total_ordering +@dataclass(frozen=True) +class TaskIdentifier: + name: str + version: str + workflow_id: UUID | None = None + + def __post_init__(self) -> None: + if self.workflow_id is not None and self.workflow_id.int == 0: + raise ValueError("Workflow ID must not be Nil") + + def __lt__(self, other: object) -> bool: + if not isinstance(other, TaskIdentifier): + return NotImplemented + return (self.name, self.version, self.workflow_id.int if self.workflow_id else 0) < ( + other.name, + other.version, + other.workflow_id.int if other.workflow_id else 0, + ) + + @classmethod + def from_message(cls, identifier: core_pb.TaskIdentifier) -> Self: + return cls( + name=identifier.name, + version=identifier.version, + workflow_id=uuid_message_to_uuid(identifier.workflow_id) if identifier.workflow_id is not None else None, + ) + def to_message(self) -> core_pb.TaskIdentifier: - return core_pb.TaskIdentifier(name=self.name, version=self.version) + return core_pb.TaskIdentifier( + name=self.name, version=self.version, workflow_id=uuid_to_uuid_message(self.workflow_id) + ) + + def unbind(self) -> UnboundTaskIdentifier: + """Return the implementation identifier without execution scope.""" + return UnboundTaskIdentifier(self.name, self.version) @classmethod def from_name_and_version(cls, name: str, version: str) -> Self: @@ -83,13 +137,8 @@ def from_name_and_version(cls, name: str, version: str) -> Self: Returns: Task metadata """ - if name == "": - raise ValueError("A task name is required") - if len(name) > 256: - raise ValueError("The task name is too long") - - _parse_version(version) # validate the version string - return cls(name=name, version=version) + identifier = UnboundTaskIdentifier.from_name_and_version(name, version) + return cls(name=identifier.name, version=identifier.version) @dataclass @@ -399,7 +448,7 @@ def to_message(self) -> workflows_pb.Path: @dataclass(order=True, frozen=True) class ReleaseContent: fingerprint: str - tasks: list[TaskIdentifier] + tasks: list[UnboundTaskIdentifier] files: list[FilesystemNode] = field(default_factory=list) runner_object_path: str = "" command_override: list[str] = field(default_factory=list) @@ -409,7 +458,7 @@ def from_message(cls, content: workflows_pb.ReleaseContent) -> "ReleaseContent": """Convert a ReleaseContent protobuf message to a ReleaseContent object.""" return cls( fingerprint=content.fingerprint, - tasks=[TaskIdentifier.from_message(task) for task in content.tasks], + tasks=[UnboundTaskIdentifier.from_message(task) for task in content.tasks], files=[FilesystemNode.from_message(path) for path in content.files], runner_object_path=content.runner_object_path, command_override=list(content.command_override), @@ -419,7 +468,7 @@ def to_message(self) -> workflows_pb.ReleaseContent: """Convert a ReleaseContent object to a ReleaseContent protobuf message.""" return workflows_pb.ReleaseContent( fingerprint=self.fingerprint, - tasks=[core_pb.UnboundTaskIdentifier(name=task.name, version=task.version) for task in self.tasks], + tasks=[task.to_message() for task in self.tasks], files=[path.to_message() for path in self.files], runner_object_path=self.runner_object_path, command_override=self.command_override, @@ -454,20 +503,52 @@ def to_message(self) -> workflows_pb.WorkflowPresentation: ) -@dataclass(order=True, frozen=True) +@dataclass(frozen=True) class Workflow: - slug: str + reference: str name: str description: str releases: list["WorkflowRelease"] = field(default_factory=list) summary: str = "" presentation: WorkflowPresentation | None = field(default=None, compare=False) + id: UUID | None = None + project_slug: str = "" + legacy_task_aliases: bool | None = None + + @property + def slug(self) -> str: + """Compatibility alias for the remote workflow reference.""" + return self.reference + + def __lt__(self, other: object) -> bool: + if not isinstance(other, Workflow): + return NotImplemented + return (self.reference, self.name, self.id.int if self.id is not None else -1) < ( + other.reference, + other.name, + other.id.int if other.id is not None else -1, + ) + + def __le__(self, other: object) -> bool: + if not isinstance(other, Workflow): + return NotImplemented + return not other < self + + def __gt__(self, other: object) -> bool: + if not isinstance(other, Workflow): + return NotImplemented + return other < self + + def __ge__(self, other: object) -> bool: + if not isinstance(other, Workflow): + return NotImplemented + return not self < other @classmethod def from_message(cls, workflow: workflows_pb.Workflow) -> "Workflow": """Convert a Workflow protobuf message to a Workflow object.""" return cls( - slug=workflow.reference, + reference=workflow.reference, name=workflow.name, description=workflow.description, releases=[WorkflowRelease.from_message(release) for release in workflow.releases], @@ -475,17 +556,23 @@ def from_message(cls, workflow: workflows_pb.Workflow) -> "Workflow": presentation=WorkflowPresentation.from_message(workflow.presentation) if workflow.presentation is not None else None, + id=uuid_message_to_uuid(workflow.id) if workflow.id is not None else None, + project_slug=workflow.project_slug, + legacy_task_aliases=workflow.legacy_task_aliases if workflow.has_field("legacy_task_aliases") else None, ) def to_message(self) -> workflows_pb.Workflow: """Convert a Workflow object to a Workflow protobuf message.""" return workflows_pb.Workflow( - reference=self.slug, + reference=self.reference, name=self.name, description=self.description, releases=[release.to_message() for release in self.releases], summary=self.summary, presentation=self.presentation.to_message() if self.presentation is not None else None, + id=uuid_to_uuid_message(self.id), + project_slug=self.project_slug, + legacy_task_aliases=self.legacy_task_aliases, ) @@ -546,7 +633,7 @@ def from_undeploy_message( @dataclass class NextTaskToRun: cluster_slug: str - identifiers: dict[TaskIdentifier, type[Any]] + identifiers: dict[UnboundTaskIdentifier, type[Any]] # from message not needed, as we never return this from the server @@ -554,7 +641,7 @@ def to_message(self) -> task_pb.NextTaskToRun: """Convert a NextTaskToRun object to a NextTaskToRun protobuf message.""" return task_pb.NextTaskToRun( cluster_slug=self.cluster_slug, - identifiers=[identifier.to_message() for identifier in self.identifiers], + identifiers=[identifier.bind().to_message() for identifier in self.identifiers], ) @@ -1486,6 +1573,7 @@ class QueryFilters: name: str | None task_states: list[TaskState] cluster_slugs: list[str] = field(default_factory=list) + workflow_references: list[str] = field(default_factory=list) @classmethod def from_message(cls, filters: job_pb.QueryFilters) -> "QueryFilters": @@ -1499,6 +1587,7 @@ def from_message(cls, filters: job_pb.QueryFilters) -> "QueryFilters": name=filters.name or None, task_states=[_TASK_STATES[state] for state in filters.task_states], cluster_slugs=list(filters.cluster_slugs), + workflow_references=list(filters.workflow_references), ) def to_message(self) -> job_pb.QueryFilters: @@ -1514,4 +1603,5 @@ def to_message(self) -> job_pb.QueryFilters: if self.task_states else None, cluster_slugs=self.cluster_slugs or None, + workflow_references=self.workflow_references or None, ) diff --git a/tilebox-workflows/tilebox/workflows/jobs/client.py b/tilebox-workflows/tilebox/workflows/jobs/client.py index 7d4fb55..944a279 100644 --- a/tilebox-workflows/tilebox/workflows/jobs/client.py +++ b/tilebox-workflows/tilebox/workflows/jobs/client.py @@ -38,6 +38,8 @@ ) from tilebox.workflows.v1.telemetry_connect import TelemetryQueryServiceClientSync from tilebox.workflows.v1.telemetry_pb import LogQueryFilters, QueryJobLogsRequest, QueryJobSpansRequest +from tilebox.workflows.v1.workflows_connect import WorkflowsServiceClientSync +from tilebox.workflows.workflows.client import WorkflowReferenceLike, resolve_workflow_id, to_workflow_slug try: from IPython.display import HTML, display @@ -55,7 +57,11 @@ def display(*_args: Any, **_kwargs: Any) -> None: class JobClient: def __init__( - self, service: JobServiceClientSync, telemetry_service: TelemetryQueryServiceClientSync, tracer: WorkflowTracer + self, + service: JobServiceClientSync, + telemetry_service: TelemetryQueryServiceClientSync, + tracer: WorkflowTracer, + workflow_service: WorkflowsServiceClientSync | None = None, ) -> None: """Create a new job client. @@ -63,10 +69,12 @@ def __init__( service: The native client to use for job operations. telemetry_service: The native client to use for job telemetry queries. tracer: The tracer to use for tracing. + workflow_service: The native client used to resolve workflow references for submissions. """ self._service = service self._telemetry_service = telemetry_service self._tracer = tracer + self._workflow_service = workflow_service def submit( self, @@ -74,6 +82,8 @@ def submit( root_task_or_tasks: TaskInstance | list[TaskInstance], cluster: ClusterSlugLike | list[ClusterSlugLike] | None = None, max_retries: int = 0, + *, + workflow: WorkflowReferenceLike | None = None, ) -> Job: """Submit a new job with the given root task(s). @@ -84,6 +94,8 @@ def submit( of clusters can be provided as well, specifying the cluster for each root task. If not provided, the default cluster will be used. max_retries: The maximum number of retries for the root task(s) in case of failure. Defaults to 0. + workflow: Workflow reference or object to scope all roots to. Omission submits unscoped tasks, even when + called inside a scoped task. A failed reference lookup never falls back to unscoped submission. Returns: The job that was submitted. @@ -102,8 +114,10 @@ def submit( f"or exactly one cluster per task. But got {len(tasks)} tasks and {len(slugs)} clusters." ) + workflow_id = resolve_workflow_id(workflow, self._workflow_service) task_submissions = [ - FutureTask(i, task, [], slugs[i], max_retries, optional=False) for i, task in enumerate(tasks) + FutureTask(i, task, [], slugs[i], max_retries, optional=False, workflow_id=workflow_id) + for i, task in enumerate(tasks) ] submissions_merged = merge_future_tasks_to_submissions(task_submissions, default_cluster) if submissions_merged is None: @@ -282,6 +296,8 @@ def query( # noqa: C901, PLR0913, PLR0917 name: str | None = None, task_states: TaskState | list[TaskState] | None = None, clusters: ClusterSlugLike | list[ClusterSlugLike] | None = None, + *, + workflows: WorkflowReferenceLike | list[WorkflowReferenceLike] | None = None, ) -> list[Job]: """List jobs matching the given filters. @@ -308,6 +324,8 @@ def query( # noqa: C901, PLR0913, PLR0917 least one task in any of the selected states are returned. clusters: A cluster or list of clusters to filter jobs by. If specified, only jobs submitted to one of the selected clusters are returned. + workflows: Workflow references or objects. Matches jobs containing any task scoped to any selected + workflow. Unscoped tasks do not match; omission adds no workflow restriction. Returns: A list of jobs matching the given filters. @@ -357,6 +375,7 @@ def query( # noqa: C901, PLR0913, PLR0917 task_states = [task_states] cluster_slugs = _to_cluster_filter_slugs(clusters) + selected_workflows = workflows if isinstance(workflows, list) else [workflows] if workflows is not None else [] filters = QueryFilters( time_interval=time_interval, @@ -366,6 +385,7 @@ def query( # noqa: C901, PLR0913, PLR0917 name=name, task_states=task_states, cluster_slugs=cluster_slugs, + workflow_references=[to_workflow_slug(workflow) for workflow in selected_workflows], ) def request(page: PaginationProtocol) -> QueryJobsResponse: diff --git a/tilebox-workflows/tilebox/workflows/runner/executor.py b/tilebox-workflows/tilebox/workflows/runner/executor.py index db4f740..343ec0f 100644 --- a/tilebox-workflows/tilebox/workflows/runner/executor.py +++ b/tilebox-workflows/tilebox/workflows/runner/executor.py @@ -71,7 +71,7 @@ def execute_task( return FailedTask(task.id, task.display, was_workflow_error=False, progress_updates=[]) try: - task_class = self.runner.tasks_by_identifier[task.identifier] + task_class = self.runner.tasks_by_identifier[task.identifier.unbind()] except KeyError: self.task_logger.error( f"Task {task.id} has unknown identifier {task.identifier}.", @@ -181,6 +181,7 @@ def submit_subtask( cluster=cluster, max_retries=max_retries, optional=optional, + workflow_id=self.current_task.identifier.workflow_id, ) self._sub_tasks.append(subtask) return subtask diff --git a/tilebox-workflows/tilebox/workflows/runner/runner.py b/tilebox-workflows/tilebox/workflows/runner/runner.py index 579d5e7..fd31772 100644 --- a/tilebox-workflows/tilebox/workflows/runner/runner.py +++ b/tilebox-workflows/tilebox/workflows/runner/runner.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING from tilebox.workflows.cache import JobCache -from tilebox.workflows.data import RunnerContext, TaskIdentifier +from tilebox.workflows.data import RunnerContext, UnboundTaskIdentifier from tilebox.workflows.task import Task, TaskMeta if TYPE_CHECKING: @@ -28,7 +28,7 @@ def __init__( ) -> None: self.cache = cache self.context = context - self._tasks_by_identifier: dict[TaskIdentifier, type[Task]] = {} + self._tasks_by_identifier: dict[UnboundTaskIdentifier, type[Task]] = {} for task in tasks or []: self.register(task) @@ -52,11 +52,11 @@ def register(self, task: type[Task]) -> None: self._tasks_by_identifier[meta.identifier] = task @property - def task_identifiers(self) -> list[TaskIdentifier]: + def task_identifiers(self) -> list[UnboundTaskIdentifier]: return list(self._tasks_by_identifier) @property - def tasks_by_identifier(self) -> dict[TaskIdentifier, type[Task]]: + def tasks_by_identifier(self) -> dict[UnboundTaskIdentifier, type[Task]]: return self._tasks_by_identifier def connect_to(self, client: Client, cluster: ClusterSlugLike | None = None) -> TaskRunner: diff --git a/tilebox-workflows/tilebox/workflows/runner/worker_service.py b/tilebox-workflows/tilebox/workflows/runner/worker_service.py index bcb975b..b4b1cd2 100644 --- a/tilebox-workflows/tilebox/workflows/runner/worker_service.py +++ b/tilebox-workflows/tilebox/workflows/runner/worker_service.py @@ -92,10 +92,7 @@ async def list_registered_tasks(self, request: Empty, context: RequestContext) - logger.debug("ListRegisteredTasks RPC called") self._authorize(context) runner = self._loaded_runner(context) - identifiers = [ - core_pb.UnboundTaskIdentifier(name=identifier.name, version=identifier.version) - for identifier in runner.task_identifiers - ] + identifiers = [identifier.to_message() for identifier in runner.task_identifiers] logger.debug(f"ListRegisteredTasks RPC returning {len(identifiers)} task identifier(s)") return core_pb.UnboundTaskIdentifiers(identifiers=identifiers) diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py index 2403b90..4e9bcc2 100644 --- a/tilebox-workflows/tilebox/workflows/task.py +++ b/tilebox-workflows/tilebox/workflows/task.py @@ -12,7 +12,13 @@ from protobuf import Message from tilebox.workflows._serialization import decode_json, encode_json_field, encode_json_fields -from tilebox.workflows.data import RunnerContext, TaskIdentifier, TaskSubmissionGroup, TaskSubmissions +from tilebox.workflows.data import ( + RunnerContext, + TaskIdentifier, + TaskSubmissionGroup, + TaskSubmissions, + UnboundTaskIdentifier, +) if TYPE_CHECKING: from tilebox.workflows.cache import JobCache @@ -172,7 +178,7 @@ def _is_valid_execute_return_annotation(annotation: Any) -> bool: @dataclass class TaskMeta: - identifier: TaskIdentifier + identifier: UnboundTaskIdentifier executable: bool @staticmethod @@ -208,7 +214,7 @@ def _task_meta(task_or_task_execution: type | Task) -> TaskMeta: return meta -def _get_task_identifier(task_class: type) -> TaskIdentifier: +def _get_task_identifier(task_class: type) -> UnboundTaskIdentifier: """Get the task identifier for a task class. Invokes a user-defined identifier method if it exists, or generates default identifier otherwise. @@ -238,7 +244,7 @@ def identifier() -> tuple[str, str]: name = class_name version = "v0.0" - return TaskIdentifier.from_name_and_version(name, version) + return UnboundTaskIdentifier.from_name_and_version(name, version) @dataclass @@ -251,9 +257,10 @@ class FutureTask: cluster: str | None max_retries: int optional: bool + workflow_id: UUID | None = None def identifier(self) -> TaskIdentifier: - return _task_meta(self.task).identifier + return _task_meta(self.task).identifier.bind(self.workflow_id) def input(self) -> bytes: return self.task._serialize() # noqa: SLF001 @@ -391,6 +398,9 @@ class CurrentTask(Protocol): display: str | None + @property + def identifier(self) -> TaskIdentifier: ... + @property def job(self) -> CurrentJob: ... @@ -412,6 +422,8 @@ def submit_subtask( ) -> FutureTask: """Submit a subtask of the current task. + The runner inherits the executing parent's workflow scope, including unscoped legacy-alias execution. + Args: task: The subtask to submit. depends_on: List of other subtasks, previously submitted within the same task context, that this subtask diff --git a/tilebox-workflows/tilebox/workflows/workflows/client.py b/tilebox-workflows/tilebox/workflows/workflows/client.py index b26e509..55cfe7a 100644 --- a/tilebox-workflows/tilebox/workflows/workflows/client.py +++ b/tilebox-workflows/tilebox/workflows/workflows/client.py @@ -16,7 +16,8 @@ UpdateWorkflowRequest, ) -WorkflowSlugLike: TypeAlias = Workflow | str +WorkflowReferenceLike: TypeAlias = Workflow | str +WorkflowSlugLike: TypeAlias = WorkflowReferenceLike WorkflowReleaseLike: TypeAlias = WorkflowRelease | UUID | str @@ -29,19 +30,35 @@ def __init__(self, service: WorkflowsServiceClientSync) -> None: """ self._service = service - def create(self, name: str, summary: str, description: str = "") -> Workflow: - """Create a new workflow with the given name. + def create( + self, + name: str, + summary: str, + description: str = "", + *, + project_slug: str = "", + ) -> Workflow: + """Create a new scoped-only workflow with the given name. Args: name: The name of the workflow to create. summary: The short summary of the workflow to create. description: The description of the workflow to create. Defaults to an empty string. + project_slug: Portable project name, not a remote workflow reference. Returns: The created workflow. """ return Workflow.from_message( - self._service.create_workflow(CreateWorkflowRequest(name=name, summary=summary, description=description)) + self._service.create_workflow( + CreateWorkflowRequest( + name=name, + summary=summary, + description=description, + project_slug=project_slug, + legacy_task_aliases=False, + ) + ) ) def all(self) -> list[Workflow]: @@ -159,7 +176,23 @@ def undeploy_release( def to_workflow_slug(workflow: WorkflowSlugLike) -> str: - return workflow.slug if isinstance(workflow, Workflow) else workflow + return workflow.reference if isinstance(workflow, Workflow) else workflow + + +def resolve_workflow_id( + workflow: WorkflowReferenceLike | None, service: WorkflowsServiceClientSync | None +) -> UUID | None: + if workflow is None: + return None + if isinstance(workflow, str): + if not workflow: + raise ValueError("A workflow reference is required") + if service is None: + raise ValueError("A workflow service is required to resolve workflow references") + workflow = Workflow.from_message(service.get_workflow(GetWorkflowRequest(workflow_reference=workflow))) + if workflow.id is None or workflow.id.int == 0: + raise ValueError("Workflow must have a non-Nil ID to scope tasks") + return workflow.id def release_id(release_or_id: WorkflowReleaseLike) -> UUID: