Add additional fields to JobDescription (#2204)

* Add additional fields to JobDescription

* Hide internal fields from public api

neuro-sdk/docs/jobs_reference.rst

@@ -698,6 +698,11 @@ JobDescription
       cannot be scheduled because the lack of computation
       cluster resources (memory, CPU/GPU etc), :class:`float`
 
+   .. attribute:: _internal
+
+      Some internal info about job used by platform. Should not be used.
+
+
 JobRestartPolicy
 ================
 

neuro-sdk/src/neuro_sdk/jobs.py

@@ -197,6 +197,13 @@ def __repr__(self) -> str:
         return repr(self.value)
 
 
+@dataclass(frozen=True)
+class JobDescriptionInternal:
+    materialized: bool = False
+    being_dropped: bool = False
+    logs_removed: bool = False
+
+
 @dataclass(frozen=True)
 class JobDescription:
     id: str
@@ -220,6 +227,7 @@ class JobDescription:
     preset_name: Optional[str] = None
     preemptible_node: bool = False
     privileged: bool = False
+    _internal: JobDescriptionInternal = JobDescriptionInternal()
 
 
 @dataclass(frozen=True)
@@ -404,6 +412,9 @@ async def list(
         reverse: bool = False,
         limit: Optional[int] = None,
         cluster_name: Optional[str] = None,
+        _materialized: Optional[bool] = None,
+        _being_dropped: Optional[bool] = False,
+        _logs_removed: Optional[bool] = False,
     ) -> AsyncIterator[JobDescription]:
         url = self._config.api_url / "jobs"
         headers = {"Accept": "application/x-ndjson"}
@@ -432,6 +443,12 @@ async def list(
             params.add("reverse", "1")
         if limit is not None:
             params.add("limit", str(limit))
+        if _materialized is not None:
+            params.add("materialized", str(_materialized))
+        if _being_dropped is not None:
+            params.add("being_dropped", str(_being_dropped))
+        if _logs_removed is not None:
+            params.add("logs_removed", str(_logs_removed))
         auth = await self._config._api_auth()
         async with self._core.request(
             "GET", url, headers=headers, params=params, auth=auth
@@ -1005,6 +1022,11 @@ def _job_description_from_api(res: Dict[str, Any], parse: Parser) -> JobDescript
         life_span=life_span,
         schedule_timeout=res.get("schedule_timeout", None),
         preset_name=res.get("preset_name"),
+        _internal=JobDescriptionInternal(
+            materialized=res.get("materialized", False),
+            being_dropped=res.get("being_dropped", False),
+            logs_removed=res.get("logs_removed", False),
+        ),
     )
 
 

neuro-sdk/tests/test_jobs.py

@@ -518,6 +518,71 @@ async def handler(request: web.Request) -> web.Response:
         assert ret.preemptible_node
 
 
+async def test_status_being_dropped(
+    aiohttp_server: _TestServerFactory, make_client: _MakeClient
+) -> None:
+    JSON = {
+        "status": "failed",
+        "id": "job-id",
+        "description": "This is job description, not a history description",
+        "http_url": "http://my_host:8889",
+        "history": {
+            "created_at": "2018-08-29T12:23:13.981621+00:00",
+            "started_at": "2018-08-29T12:23:15.988054+00:00",
+            "finished_at": "2018-08-29T12:59:31.427795+00:00",
+            "reason": "ContainerCannotRun",
+            "description": "Not enough coffee",
+        },
+        "scheduler_enabled": True,
+        "preemptible_node": True,
+        "pass_config": False,
+        "being_dropped": True,
+        "logs_removed": True,
+        "owner": "owner",
+        "cluster_name": "default",
+        "uri": "job://default/owner/job-id",
+        "container": {
+            "image": "submit-image-name",
+            "command": "submit-command",
+            "http": {"port": 8181},
+            "resources": {
+                "memory_mb": "4096",
+                "cpu": 7.0,
+                "shm": True,
+                "gpu": 1,
+                "gpu_model": "test-gpu-model",
+            },
+            "volumes": [
+                {
+                    "src_storage_uri": "storage://test-user/path_read_only",
+                    "dst_path": "/container/read_only",
+                    "read_only": True,
+                },
+                {
+                    "src_storage_uri": "storage://test-user/path_read_write",
+                    "dst_path": "/container/path_read_write",
+                    "read_only": False,
+                },
+            ],
+        },
+    }
+
+    async def handler(request: web.Request) -> web.Response:
+        return web.json_response(JSON)
+
+    app = web.Application()
+    app.router.add_get("/jobs/job-id", handler)
+
+    srv = await aiohttp_server(app)
+
+    async with make_client(srv.make_url("/")) as client:
+        ret = await client.jobs.status("job-id")
+
+        assert ret == _job_description_from_api(JSON, client.parse)
+        assert ret._internal.being_dropped
+        assert ret._internal.logs_removed
+
+
 async def test_status_with_http(
     aiohttp_server: _TestServerFactory, make_client: _MakeClient
 ) -> None: