datatorch 0.6.2__tar.gz → 0.7.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datatorch
-Version: 0.6.2
+Version: 0.7.1
 Summary: A CLI and library for interacting with DataTorch.
 Home-page: https://github.com/datatorch/python
 Author: DataTorch
@@ -95,4 +95,4 @@ 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
+source env/bin/activate && pip3 install dist/datatorch-0.6.3-py3-none-any.whl && pip3 install pycocotools

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

@@ -45,4 +45,4 @@ 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
+source env/bin/activate && pip3 install dist/datatorch-0.6.3-py3-none-any.whl && pip3 install pycocotools

{datatorch-0.6.2 → datatorch-0.7.1}/datatorch/agent/__init__.py

@@ -12,8 +12,7 @@ from .agent import Agent, tasks
 from gql import Client as GqlClient
 from gql.transport.exceptions import TransportClosed, TransportServerError
 from gql.transport.websockets import WebsocketsTransport
-from websockets.exceptions import InvalidMessage, InvalidURI
-from websockets import ConnectionClosedError
+from websockets.exceptions import ConnectionClosed, InvalidMessage, InvalidURI
 
 from datatorch.api import Client as DtClient
 from datatorch.utils.package import get_version
@@ -110,12 +109,17 @@ async def start() -> None:
                 await Agent.run(session)
 
         except (
-            ConnectionClosedError,
-            ConnectionRefusedError,
+            # ConnectionClosed covers both the error and the "OK" close
+            # (a server-side restart closes cleanly); OSError covers the
+            # socket-level failures seen right after a laptop wakes
+            # (DNS gaierror, ConnectionReset, ConnectionRefused, timeouts).
+            ConnectionClosed,
+            OSError,
             IncompleteReadError,
             TransportServerError,
             TransportClosed,
             InvalidMessage,
+            asyncio.TimeoutError,
         ) as e:
             await _exit_jobs()
             await _close_transport(transport)
@@ -139,12 +143,12 @@ async def stop() -> None:
 
     print(" ")
 
-    logger.warning("Gracefully exiting agent.")
+    logger.warning("Gracefully exiting runner.")
 
-    logger.info("Closing agent jobs.")
+    logger.info("Closing runner jobs.")
     await _exit_jobs()
 
-    logger.info("Closing agent jobs.")
+    logger.info("Closing runner jobs.")
     await _exit_jobs()
 
     logger.info("Closing all other tasks.")

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

@@ -4,6 +4,9 @@ from .directory import agent_directory
 from gql.client import AsyncClientSession
 
 from typing_extensions import TypedDict
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class Log(TypedDict):
@@ -236,6 +239,17 @@ class AgentApiClient(object):
         # fmt: on
         return self.execute(mutate, params={"id": step_id, "logs": logs})
 
+    async def force_reconnect(self) -> None:
+        """Closes the websocket so the steps subscription raises and the
+        outer loop in `agent.start()` reconnects with backoff."""
+        transport = getattr(self.session, "transport", None)
+        if transport is None or getattr(transport, "websocket", None) is None:
+            return
+        try:
+            await transport.close()
+        except Exception as e:
+            logger.debug(f"Transport close raised while forcing reconnect: {e}")
+
     async def execute(self, query, *args, params: dict = {}, **kwargs) -> dict:
         """Wrapper around execute"""
         removed_none = dict((k, v) for k, v in params.items() if v is not None)

{datatorch-0.6.2 → datatorch-0.7.1}/datatorch/agent/monitoring.py

@@ -67,10 +67,17 @@ class AgentSystemStats(object):
         print_stats(stats)
         return stats
 
+    # Consecutive failed reports before the connection is considered dead.
+    # The steps subscription can sit silently on a half-open socket (e.g.
+    # after sleep); the report failing is the signal that reaches us.
+    MAX_REPORT_FAILURES = 3
+    REPORT_TIMEOUT = 30
+
     def __init__(self, agent, sample_rate=60):
         self.agent = agent
         self.sample_rate = sample_rate
         self.sample = 0
+        self.report_failures = 0
 
     async def start(self):
         try:
@@ -90,7 +97,25 @@ class AgentSystemStats(object):
         while True:
             self.sample += 1
             stats = self.stats()
-            loop = asyncio.get_running_loop()
             if self.sample != 1:
-                loop.create_task(self.agent.api.metrics(stats))
+                await self._report(stats)
             await asyncio.sleep(self.sample_rate)
+
+    async def _report(self, stats):
+        """Sends one metric report; repeated failures force a reconnect."""
+        try:
+            await asyncio.wait_for(
+                self.agent.api.metrics(stats), timeout=self.REPORT_TIMEOUT
+            )
+            self.report_failures = 0
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            self.report_failures += 1
+            logger.warning(
+                f"Metric report failed ({self.report_failures}/{self.MAX_REPORT_FAILURES}): {e}"
+            )
+            if self.report_failures >= self.MAX_REPORT_FAILURES:
+                logger.error("Connection looks dead; forcing a reconnect.")
+                self.report_failures = 0
+                await self.agent.api.force_reconnect()

{datatorch-0.6.2 → datatorch-0.7.1}/datatorch/agent/pipelines/runner/runner.py

@@ -14,6 +14,12 @@ class ProcessCodeError(Exception):
     pass
 
 
+# Max size of one stdout line from an action subprocess. Output values
+# arrive as one `::var::<json>` line, so this bounds the largest single
+# output value (e.g. an http-request body), not ordinary log lines.
+STDOUT_LINE_LIMIT = 16 * 1024 * 1024
+
+
 class Runner(object):
     def __init__(self, config: dict, action: "Action"):
         self.config = config
@@ -110,6 +116,11 @@ class Runner(object):
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.STDOUT,
             env={**os.environ, **(env or {})},
+            # Outputs are reported as single `::var::<json>` stdout lines
+            # (see check_for_output), so a line can carry an entire fetched
+            # page body — asyncio's 64KB default readline limit kills the
+            # stream with LimitOverrunError on anything larger.
+            limit=STDOUT_LINE_LIMIT,
         )
         if wait:
             await process.wait()

{datatorch-0.6.2 → datatorch-0.7.1}/datatorch/api/api.py

@@ -56,6 +56,14 @@ _FILE = File.add_fragment("""
     }
     """)
 
+_DATASET = Dataset.add_fragment("""
+    query GetDataset($id: ID!) {
+      dataset: datasetById(id: $id) {
+        ...DatasetFields
+      }
+    }
+    """)
+
 _VIEWER = User.add_fragment("""
     query GetViewer {
       viewer {
@@ -119,6 +127,13 @@ class ApiClient(Client):
             File, self.query_to_class(File, _FILE, path="file", params={"fileId": id})
         )
 
+    def dataset(self, id: str) -> Dataset:
+        """Retrieve a dataset (or branch) by ID"""
+        return cast(
+            Dataset,
+            self.query_to_class(Dataset, _DATASET, path="dataset", params={"id": id}),
+        )
+
     def upload_to_default_filesource(
         self,
         project: Project,

{datatorch-0.6.2 → datatorch-0.7.1}/datatorch/api/client.py

@@ -162,7 +162,8 @@ class Client(object):
         dl_file = os.path.join(directory, name)
         dl_file = os.path.abspath(dl_file)
 
-        if os.path.isfile(dl_file) & skip == True:
+        # `isfile(x) & skip == True` only worked by bitwise-bool accident
+        if skip and os.path.isfile(dl_file):
             return dl_file, None
 
         query_string = urlencode({"download": "true", "stream": "true"})
@@ -172,8 +173,10 @@ class Client(object):
         result = requests.get(
             download_url, headers={self.token_header: self._api_token}, stream=True
         )
+        # Never write an error body to disk as a "downloaded file"
+        result.raise_for_status()
 
-        content = result.headers["content-disposition"]
+        content = result.headers.get("content-disposition", "")
         # _, value = cgi.parse_header(content)
         parser = HeaderParser()
         headers = parser.parsestr(f"Content-Disposition: {content}")

datatorch-0.7.1/datatorch/api/entity/dataset.py

@@ -0,0 +1,103 @@
+from .base import BaseEntity
+
+__all__ = "Dataset"
+
+
+_FORK_DATASET = """
+    mutation ForkDataset(
+      $datasetId: ID!
+      $branchName: String!
+      $versionId: ID
+    ) {
+      fork: forkDataset(
+        datasetId: $datasetId
+        branchName: $branchName
+        versionId: $versionId
+      ) {
+        dataset {
+          id
+          name
+          description
+          projectId
+          branchName
+          parentDatasetId
+          forkedAt
+        }
+      }
+    }
+"""
+
+
+_CREATE_DATASET = """
+    mutation CreateDataset(
+      $projectId: ID!
+      $name: String!
+      $description: String
+    ) {
+      dataset: createDataset(
+        input: {
+          projectId: $projectId
+          name: $name
+          description: $description
+        }
+      ) {
+        id
+      }
+    }
+"""
+
+
+class Dataset(BaseEntity):
+    id: str
+    name: str
+    description: str
+    project_id: str
+    kilobytes: int
+    formatted_bytes: int
+    created_at: str
+    updated_at: str
+
+    def create(self, client=None):
+        super().create(client=client)
+
+        assert self.project_id is not None
+        results = self.client.execute(
+            _CREATE_DATASET,
+            params={
+                "projectId": self.project_id,
+                "name": self.name,
+                "description": self.description,
+            },
+        )
+
+        self.id = results.get("dataset").get("id")
+
+    def fork(self, branch_name: str, version_id: str = None, client=None):
+        """
+        Forks this dataset into a branch and returns it as a new Dataset.
+
+        The branch shares the parent's storage objects and copies its
+        annotations. Pass `version_id` to start the branch from a version's
+        reconstructed state instead of the live state. The returned dataset
+        can be used anywhere a dataset id is accepted (e.g. uploads).
+
+        Note: branch attributes (`branch_name`, `parent_dataset_id`,
+        `forked_at`) are set dynamically on the returned instance; they are
+        deliberately kept out of the class annotations so the generated
+        `DatasetFields` fragment stays compatible with servers that predate
+        branching.
+        """
+        assert self.id is not None
+        if client:
+            self.client = client
+        assert self.client is not None
+
+        results = self.client.execute(
+            _FORK_DATASET,
+            params={
+                "datasetId": self.id,
+                "branchName": branch_name,
+                "versionId": version_id,
+            },
+        )
+        return Dataset(results.get("fork").get("dataset"), client=self.client)