datatorch 0.4.9.2__tar.gz → 0.6.2__tar.gz

{datatorch-0.4.9.2 → datatorch-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datatorch
-Version: 0.4.9.2
+Version: 0.6.2
 Summary: A CLI and library for interacting with DataTorch.
 Home-page: https://github.com/datatorch/python
 Author: DataTorch
@@ -65,7 +65,7 @@ Dynamic: summary
 
 ## Installation
 
-*Note*: Newer versions (>=0.4.8.2) of DataTorch client use `libmagic` for more 
+_Note_: Newer versions (>=0.4.8.2) of DataTorch client use `libmagic` for more
 accurate mimetype guessing. This falls back to guessing via extension. Consider
 installing `libmagic` if you have issues with file types uploaded to DataTorch.
 
@@ -86,3 +86,13 @@ pip3 install -r requirements.txt
 
 Open your VSCode settings, by going `Code -> Preferences -> Settings`. Search
 for "python formatting provider" and select "black" from the dropdown menu.
+
+#### Testing
+
+python3 -m build
+
+python3 -m venv env
+
+rm -rf env
+
+source env/bin/activate && pip3 install dist/datatorch-0.6.1-py3-none-any.whl && pip3 install pycocotools

{datatorch-0.4.9.2 → datatorch-0.6.2}/README.md

@@ -15,7 +15,7 @@
 
 ## Installation
 
-*Note*: Newer versions (>=0.4.8.2) of DataTorch client use `libmagic` for more 
+_Note_: Newer versions (>=0.4.8.2) of DataTorch client use `libmagic` for more
 accurate mimetype guessing. This falls back to guessing via extension. Consider
 installing `libmagic` if you have issues with file types uploaded to DataTorch.
 
@@ -36,3 +36,13 @@ pip3 install -r requirements.txt
 
 Open your VSCode settings, by going `Code -> Preferences -> Settings`. Search
 for "python formatting provider" and select "black" from the dropdown menu.
+
+#### Testing
+
+python3 -m build
+
+python3 -m venv env
+
+rm -rf env
+
+source env/bin/activate && pip3 install dist/datatorch-0.6.1-py3-none-any.whl && pip3 install pycocotools

datatorch-0.6.2/datatorch/agent/agent.py

@@ -0,0 +1,147 @@
+import os
+import logging
+import asyncio
+
+from typing import Dict, List
+from .client import AgentApiClient, AgentStepDispatch
+from .pipelines.template import Variables
+from .log_handler import AgentAPIHandler
+from .monitoring import AgentSystemStats
+from .directory import agent_directory
+
+from gql.client import AsyncClientSession
+from datatorch.agent.pipelines.step import Step
+
+logger = logging.getLogger(__name__)
+tasks: List[asyncio.Task] = []
+
+
+class Agent(object):
+    """
+    Agent protocol v2: the server sequences steps. The agent receives one
+    ready step at a time — its inputs already resolved server-side —
+    claims it, executes it, and reports the result. Job-level state and
+    cross-step dataflow live entirely on the server.
+    """
+
+    @classmethod
+    async def run(cls, session):
+        agent = cls(session)
+        await agent.process_loop()
+
+    def __init__(self, session: AsyncClientSession):
+        self.api = AgentApiClient(session)
+        self.directory = agent_directory
+        # Running step tasks, keyed by stepId, so a server cancel/timeout
+        # signal can abort the right one. Entries are cleaned up when the
+        # step finishes.
+        self._step_tasks: Dict[str, asyncio.Task] = {}
+        logger.debug(f"Switch to agent directory: {self.directory.dir}")
+        os.chdir(self.directory.dir)
+
+        self._init_logger()
+        self._init_threads()
+
+    def _init_logger(self):
+        self.logger = logging.getLogger("datatorch.agent")
+        self.logger_api_handler = AgentAPIHandler(self.api)
+        self.logger.addHandler(self.logger_api_handler)
+        self.logger.debug("Agent logger has been initalized.")
+
+    def _init_threads(self):
+        self.system_stats = AgentSystemStats(self)
+        loop = asyncio.get_running_loop()
+        task = loop.create_task(self.system_stats.start())
+        tasks.append(task)
+
+    async def process_loop(self):
+        """Drains steps missed while offline, then waits for pushes.
+
+        Runs the dispatch stream alongside a cancel-signal listener: the
+        server pushes a stop signal when a running step is canceled or
+        times out, and the agent aborts that step's task.
+        """
+        loop = asyncio.get_running_loop()
+
+        pending = await self.api.pending_steps()
+        if pending:
+            logger.info(f"Recovered {len(pending)} pending step(s) on connect.")
+        for dispatch in pending:
+            self._start_step(dispatch)
+
+        logger.info("Waiting for steps.")
+        cancel_task = loop.create_task(self._cancel_loop())
+        tasks.append(cancel_task)
+        try:
+            async for request in self.api.agent_steps():
+                self._start_step(request.get("step"))
+        finally:
+            cancel_task.cancel()
+
+    async def _cancel_loop(self):
+        """Aborts a running step's task when the server signals a stop."""
+        async for request in self.api.agent_step_cancels():
+            cancel = request.get("cancel") or {}
+            step_id = cancel.get("stepId")
+            task = self._step_tasks.get(step_id)
+            if task and not task.done():
+                logger.info(f"Received stop signal for step {step_id}.")
+                task.cancel()
+            else:
+                # Unknown/finished step — the signal is best-effort and the
+                # server has already moved the step terminal.
+                logger.debug(f"Stop signal for step {step_id} ignored (not running).")
+
+    def _start_step(self, dispatch: AgentStepDispatch):
+        """Schedules a step and tracks its task for cancellation."""
+        loop = asyncio.get_running_loop()
+        step_id = dispatch.get("stepId")
+        task = loop.create_task(self._run_step(dispatch))
+        self._step_tasks[step_id] = task
+        tasks.append(task)
+
+    async def _run_step(self, dispatch: AgentStepDispatch):
+        """Claims and executes a single dispatched step."""
+        step_id = dispatch.get("stepId")
+        step_name = dispatch.get("stepName") or dispatch.get("action")
+
+        try:
+            claimed = await self.api.claim_step(step_id)
+            if not claimed:
+                # Another process (or an earlier duplicate dispatch) owns it.
+                logger.debug(f"Step {step_id} already claimed; skipping.")
+                return
+
+            logger.info(f"Step received (name: {step_name}, id: {step_id})")
+            step = Step.from_dispatch(dispatch, api=self.api)
+
+            try:
+                outputs = await step.run(Variables())
+            except asyncio.CancelledError:
+                # Either a disconnect mid-step (the reconciler fails the
+                # orphaned RUNNING step once the agent stays offline) or a
+                # server stop signal (the step is already terminal
+                # server-side). Either way the child process was killed by
+                # the runner; don't report a result — a late completion
+                # would be absorbed by the RUNNING guard anyway.
+                logger.info(f"Step {step_id} aborted.")
+                await step.upload_logs()
+                raise
+            except Exception as e:
+                logger.error(f"Step {step_id} failed: {e}")
+                step.log(f"Step failed: {e}.")
+                await step.upload_logs()
+                await self.api.complete_step(
+                    step_id, success=False, error_message=str(e)
+                )
+                return
+
+            await self.api.complete_step(
+                step_id,
+                success=True,
+                output=outputs,
+                rendered_input=step.rendered_inputs,
+            )
+            logger.info(f"Step finished (name: {step_name}, id: {step_id})")
+        finally:
+            self._step_tasks.pop(step_id, None)

{datatorch-0.4.9.2 → datatorch-0.6.2}/datatorch/agent/client.py

@@ -1,5 +1,4 @@
-from datatorch.api.api import ApiClient
-from typing import AsyncGenerator, List, cast
+from typing import AsyncGenerator, List, Optional, cast
 from gql import gql
 from .directory import agent_directory
 from gql.client import AsyncClientSession
@@ -12,95 +11,152 @@ class Log(TypedDict):
     message: str
 
 
-class AgentPipelineConfig(TypedDict):
-    id: str
-    projectId: str
-    creatorId: str
-    lastRunNumber: int
+class AgentStepDispatch(TypedDict):
+    """A ready step pushed to this agent (agent protocol v2).
 
+    ``input`` arrives server-resolved: ``steps.*`` and ``input.*``
+    references are already substituted with concrete values.
+    """
 
-class AgentStepsConfig(TypedDict):
-    id: str
-    name: str
+    stepId: str
+    agentId: str
     action: str
-    index: int
+    stepName: Optional[str]
+    stepIndex: int
+    input: dict
+    jobId: str
+    jobName: str
+    runId: str
+    runNumber: int
+    pipelineId: str
+    projectId: str
 
 
-class AgentTriggerConfig(TypedDict):
-    id: str
-    type: str
-    event: dict
+class AgentStepRequest(TypedDict):
+    step: AgentStepDispatch
 
 
-class AgentRunConfig(TypedDict):
-    id: str
-    name: str
-    text: str
-    config: dict
-    runNumber: int
-    pipeline: AgentPipelineConfig
-    trigger: AgentTriggerConfig
+class AgentStepCancel(TypedDict):
+    """A stop signal for a step this agent is currently running."""
 
+    stepId: str
+    agentId: str
 
-class AgentJobConfig(TypedDict):
-    id: str
-    name: str
-    run: AgentRunConfig
-    steps: List[AgentStepsConfig]
 
+class AgentStepCancelRequest(TypedDict):
+    cancel: AgentStepCancel
 
-class AgentRequestConfig(TypedDict):
-    job: AgentJobConfig
+
+STEP_DISPATCH_FIELDS = """
+    stepId
+    agentId
+    action
+    stepName
+    stepIndex
+    input
+    jobId
+    jobName
+    runId
+    runNumber
+    pipelineId
+    projectId
+"""
 
 
 class AgentApiClient(object):
     def __init__(self, session: AsyncClientSession):
         self.session = session
 
-    def agent_jobs(self):
-        """Subscriptions to the agent job assignment namespace."""
+    def agent_steps(self):
+        """Subscribe to ready steps dispatched to this agent."""
         # fmt: off
         sub = gql("""
             subscription {
-                job: agentJobs {
-                    id
-                    name
-                    run {
-                        id
-                        name
-                        text
-                        config
-                        runNumber
-                        pipeline {
-                            id
-                            projectId
-                            creatorId
-                            lastRunNumber
-                        }
-                        trigger {
-                            id
-                            trigger {
-                                id
-                                type
-                                config
-                            }
-                            event
-                        }
-                    }
-                    steps {
-                        id
-                        name
-                        index
-                        action
-                    }
+                step: agentSteps {
+                    %s
+                }
+            }
+        """ % STEP_DISPATCH_FIELDS)
+        # fmt: on
+        return cast(AsyncGenerator[AgentStepRequest, None], self.session.subscribe(sub))
+
+    def agent_step_cancels(self):
+        """Subscribe to stop signals for steps this agent is running.
+
+        Best effort: the server has already moved the step terminal, so a
+        missed signal only means the step runs to harmless completion.
+        """
+        # fmt: off
+        sub = gql("""
+            subscription {
+                cancel: agentStepCancels {
+                    stepId
+                    agentId
                 }
             }
         """)
         # fmt: on
         return cast(
-            AsyncGenerator[AgentRequestConfig, None], self.session.subscribe(sub)
+            AsyncGenerator[AgentStepCancelRequest, None],
+            self.session.subscribe(sub),
         )
 
+    async def pending_steps(self) -> List[AgentStepDispatch]:
+        """Steps dispatched to this agent but not yet claimed.
+
+        Called on (re)connect: pushes missed while offline are recovered
+        here.
+        """
+        # fmt: off
+        query = """
+            query PendingSteps {
+                steps: agentPendingSteps {
+                    %s
+                }
+            }
+        """ % STEP_DISPATCH_FIELDS
+        # fmt: on
+        result = await self.execute(query)
+        return cast(List[AgentStepDispatch], result.get("steps", []))
+
+    async def claim_step(self, step_id: str) -> bool:
+        """Atomically claim a dispatched step. False = already claimed."""
+        # fmt: off
+        mutate = """
+            mutation ClaimStep($id: ID!) {
+                claimed: claimPipelineStep(id: $id)
+            }
+        """
+        # fmt: on
+        result = await self.execute(mutate, params={"id": step_id})
+        return bool(result.get("claimed"))
+
+    async def complete_step(
+        self,
+        step_id: str,
+        success: bool,
+        output: Optional[dict] = None,
+        rendered_input: Optional[dict] = None,
+        error_message: Optional[str] = None,
+    ) -> bool:
+        """Report the result of a claimed step."""
+        # fmt: off
+        mutate = """
+            mutation CompleteStep($id: ID!, $input: CompletePipelineStepInput!) {
+                completed: completePipelineStep(id: $id, input: $input)
+            }
+        """
+        # fmt: on
+        step_input = {"success": success}
+        if output is not None:
+            step_input["output"] = output
+        if rendered_input is not None:
+            step_input["renderedInput"] = rendered_input
+        if error_message is not None:
+            step_input["errorMessage"] = error_message
+        result = await self.execute(mutate, params={"id": step_id, "input": step_input})
+        return bool(result.get("completed"))
+
     def initial_metrics(self, metrics):
         # fmt: off
         mutate = """
@@ -170,45 +226,6 @@ class AgentApiClient(object):
         params = {"agentId": agent_directory.settings.agent_id, **metrics}
         return self.execute(mutate, params=params)
 
-    def update_step(self, values: dict):
-        # fmt: off
-        mutate = """
-            mutation UpdateStepRun(
-                $id: ID!
-                $inputs: JSON
-                $outputs: JSON
-                $status: PipelineStepStatus
-                $startedAt: DateTime
-                $finishedAt: DateTime
-            ) {
-                updatePipelineStep(
-                    id: $id
-                    input: {
-                        status: $status
-                        output: $outputs
-                        input: $inputs
-                        startedAt: $startedAt
-                        finishedAt: $finishedAt
-                    }
-                )
-            }
-        """
-        # fmt: on
-        return self.execute(mutate, params=values)
-
-    async def update_job(self, values: dict):
-        # fmt: off
-        mutate = """
-            mutation UpdateJob($id: ID!, $status: PipelineJobStatus) {
-                updatePipelineJobRun(
-                    id: $id,
-                    input: { status: $status }
-                )
-            }
-        """
-        # fmt: on
-        return await self.execute(mutate, params=values)
-
     def upload_step_logs(self, step_id: str, logs: List[Log]):
         # fmt: off
         mutate = """
@@ -227,7 +244,3 @@ class AgentApiClient(object):
         return await self.session.execute(
             query, *args, variable_values=removed_none, **kwargs
         )
-
-
-def create_client():
-    ApiClient

{datatorch-0.4.9.2 → datatorch-0.6.2}/datatorch/agent/pipelines/action/action.py

@@ -117,10 +117,6 @@ class Action(object):
                         variable_value = json.loads(variable_value)
                     variables.add_input(k, variable_value)
 
-        if self.step is not None:
-            # Update steps output after casting.
-            await self.step.update(inputs=variables.inputs)
-
         logger.debug(f"Inputs for '{self.full_name}': {json.dumps(variables.inputs)}")
 
         output = self.get_cached(variables)

datatorch-0.6.2/datatorch/agent/pipelines/job/job.py

@@ -0,0 +1,53 @@
+import logging
+
+from ..step import Step
+from ..template import Variables
+from ..resolver import UnresolvedReferenceError, resolve_step_input
+
+logger = logging.getLogger("datatorch.agent.job")
+
+
+class Job(object):
+    """
+    LOCAL-MODE step sequencer (``datatorch pipeline run <yaml>``).
+
+    Against the server, jobs are sequenced by the server-side orchestrator
+    and this class is not involved — the agent executes one dispatched
+    step at a time. Local mode replicates the server's behavior so the
+    same yaml works in both places: steps run sequentially, each step's
+    ``${{ steps.<name>.output.<key> }}`` / ``${{ input.<key> }}``
+    references are resolved (strictly) from earlier outputs before it
+    runs, and the first failure skips the rest.
+    """
+
+    def __init__(self, config: dict, trigger_input: dict = None):
+        self.config = config
+        self.name = config.get("name")
+        self.trigger_input = trigger_input or {}
+
+    async def run(self):
+        """Runs each step of the job, resolving references in between."""
+        outputs_by_name = {}
+
+        for step_config in self.config.get("steps", []):
+            name = step_config.get("name")
+            try:
+                resolved, unresolved = resolve_step_input(
+                    step_config.get("inputs", {}),
+                    outputs_by_name,
+                    self.trigger_input,
+                )
+                if unresolved:
+                    raise UnresolvedReferenceError(unresolved)
+
+                step = Step.from_dict({**step_config, "inputs": resolved})
+                outputs = await step.run(Variables())
+                if name:
+                    outputs_by_name[name] = outputs or {}
+            except Exception as e:
+                logger.error(f"Job {self.name} failed on step '{name}': {e}")
+                logger.error("Skipping remaining steps.")
+                raise
+
+        logger.info(f"Successfully completed job '{self.name}'.")
+        return outputs_by_name

datatorch-0.6.2/datatorch/agent/pipelines/resolver.py

@@ -0,0 +1,156 @@
+"""
+Local-mode port of the server's step input resolver.
+
+When a pipeline runs against the DataTorch server, the server resolves
+``${{ steps.<name>.output.<key> }}``, ``${{ input.<key> }}`` and
+``${{ jobs.<name>.outputs.<key> }}`` references before dispatching each
+step. Local mode (``datatorch pipeline run``) replicates those exact
+semantics so a pipeline yaml behaves identically in both places:
+
+- A reference spanning the whole string resolves to the raw value
+  (numbers/objects preserved); embedded references interpolate as strings.
+- ``outputs``/``inputs`` are accepted as aliases of ``output``/``input``.
+- Unresolvable references are an error (strict — matching the server,
+  which fails the step at dispatch).
+
+Cross-job note: ``${{ jobs.<name>.outputs.<key> }}`` is recognized as a
+server namespace here for parity, but local mode runs a single job at a
+time (there is no cross-job sequencer — see ``job.Job``), so no job
+outputs are supplied and such a reference resolves as unresolvable, i.e.
+raises rather than passing a literal ``${{ }}`` to the action. A yaml that
+uses job dependencies is a server-executed workflow.
+"""
+
+import json
+import re
+from typing import Any, Dict, List, Set, Tuple
+
+# Named groups so the combined alternation reads by name, not by positional
+# index that shifts when a namespace is added — mirrors the server resolver.
+# The name group is ``[^}]+?`` (not ``.+?``) so it can't backtrack across a
+# ``}}`` boundary — a string that both starts and ends with a ref must not be
+# mistaken for one spanning ref.
+STEP_REF = (
+    r"\$\{\{\s*steps\.(?P<stepName>[^}]+?)\.outputs?\.(?P<stepKey>[\w.\-]+)\s*\}\}"
+)
+INPUT_REF = r"\$\{\{\s*inputs?\.(?P<inputKey>[\w.\-]+)\s*\}\}"
+JOB_REF = r"\$\{\{\s*jobs\.(?P<jobName>[^}]+?)\.outputs?\.(?P<jobKey>[\w.\-]+)\s*\}\}"
+
+SERVER_REF = re.compile(f"{STEP_REF}|{INPUT_REF}|{JOB_REF}")
+WHOLE_STRING_REF = re.compile(f"^(?:{STEP_REF}|{INPUT_REF}|{JOB_REF})$")
+
+_MISSING = object()
+
+
+class UnresolvedReferenceError(Exception):
+    def __init__(self, refs: List[str]):
+        self.refs = refs
+        super().__init__(
+            "Unresolvable input reference(s): "
+            + ", ".join(refs)
+            + ". References must point to a named earlier step output, "
+            + "a trigger input, or a needed job output."
+        )
+
+
+def _lookup_mapping(mapping: Dict[str, dict], name, key):
+    if name is None or key is None:
+        return _MISSING
+    values = mapping.get(name.strip(), _MISSING)
+    if values is _MISSING or not isinstance(values, dict):
+        return _MISSING
+    return values.get(key, _MISSING)
+
+
+def _lookup(
+    outputs_by_name: Dict[str, dict],
+    trigger_input: dict,
+    job_outputs_by_name: Dict[str, dict],
+    groups: Dict[str, Any],
+):
+    step = _lookup_mapping(
+        outputs_by_name, groups.get("stepName"), groups.get("stepKey")
+    )
+    if groups.get("stepName") is not None:
+        return step
+    job = _lookup_mapping(
+        job_outputs_by_name, groups.get("jobName"), groups.get("jobKey")
+    )
+    if groups.get("jobName") is not None:
+        return job
+    input_key = groups.get("inputKey")
+    if input_key is not None:
+        return trigger_input.get(input_key, _MISSING)
+    return _MISSING
+
+
+def _resolve(
+    value, outputs_by_name, trigger_input, job_outputs_by_name, misses: Set[str]
+):
+    if isinstance(value, str):
+        whole = WHOLE_STRING_REF.match(value)
+        if whole:
+            resolved = _lookup(
+                outputs_by_name,
+                trigger_input,
+                job_outputs_by_name,
+                whole.groupdict(),
+            )
+            if resolved is _MISSING:
+                misses.add(value.strip())
+                return value
+            return resolved
+
+        def replace(match: "re.Match") -> str:
+            resolved = _lookup(
+                outputs_by_name,
+                trigger_input,
+                job_outputs_by_name,
+                match.groupdict(),
+            )
+            if resolved is _MISSING:
+                misses.add(match.group(0))
+                return match.group(0)
+            if isinstance(resolved, (dict, list)):
+                return json.dumps(resolved)
+            if isinstance(resolved, bool):
+                # Match the server's String(true|false)
+                return "true" if resolved else "false"
+            return str(resolved)
+
+        return SERVER_REF.sub(replace, value)
+    if isinstance(value, list):
+        return [
+            _resolve(v, outputs_by_name, trigger_input, job_outputs_by_name, misses)
+            for v in value
+        ]
+    if isinstance(value, dict):
+        return {
+            k: _resolve(v, outputs_by_name, trigger_input, job_outputs_by_name, misses)
+            for k, v in value.items()
+        }
+    return value
+
+
+def resolve_step_input(
+    inputs: Any,
+    outputs_by_name: Dict[str, dict],
+    trigger_input: dict = {},
+    job_outputs_by_name: Dict[str, dict] = {},
+) -> Tuple[Any, List[str]]:
+    """Deep-resolve server-namespace references in a step's inputs.
+
+    Returns ``(resolved, unresolved)`` where ``unresolved`` lists the
+    references (as authored) that could not be resolved. ``job_outputs_by_name``
+    is empty in local mode (single-job execution), so ``${{ jobs.* }}``
+    references resolve as unresolvable.
+    """
+    misses: Set[str] = set()
+    resolved = _resolve(
+        inputs,
+        outputs_by_name,
+        trigger_input or {},
+        job_outputs_by_name or {},
+        misses,
+    )
+    return resolved, sorted(misses)