Merge 7fdce23 into cb2a1e9

gateway-experiments · Sep 23, 2019 · f12b9d6 · f12b9d6
2 parents cb2a1e9 + 7fdce23
commit f12b9d6
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 7 deletions.
diff --git a/tests/test_application_master.py b/tests/test_application_master.py
@@ -29,7 +29,7 @@ def test_job(self, request_mock):
         request_mock.assert_called_with('/proxy/app_100500/ws/v1/mapreduce/jobs/job_100500')
 
     def test_job_attempts(self, request_mock):
-        self.app.job_attempts('app_1')
+        self.app.job_attempts('app_1', 'job_2')
 
     def test_job_counters(self, request_mock):
         self.app.job_counters('app_1', 'job_2')
@@ -57,9 +57,21 @@ def test_task_attempts(self, request_mock):
 
     def test_task_attempt(self, request_mock):
         self.app.task_attempt('app_1', 'job_2', 'task_3', 'attempt_4')
-        request_mock.assert_called_with('/proxy/app_1/ws/v1/mapreduce/jobs/job_2/tasks/task_3/attempt/attempt_4')
+        request_mock.assert_called_with('/proxy/app_1/ws/v1/mapreduce/jobs/job_2/tasks/task_3/attempts/attempt_4')
+
+    def test_task_attempt_state(self, request_mock):
+        self.app.task_attempt_state('app_1', 'job_2', 'task_3', 'attempt_4')
+        request_mock.assert_called_with('/proxy/app_1/ws/v1/mapreduce/jobs/job_2/tasks/task_3/attempts/attempt_4/state')
+
+    def test_task_attempt_state_kill(self, request_mock):
+        self.app.task_attempt_state_kill('app_1', 'job_2', 'task_3', 'attempt_4')
+        request_mock.assert_called_with(
+            '/proxy/app_1/ws/v1/mapreduce/jobs/job_2/tasks/task_3/attempts/attempt_4/state',
+            'PUT', data={'state': 'KILLED'}
+        )
 
     def test_task_attempt_counters(self, request_mock):
         self.app.task_attempt_counters('app_1', 'job_2', 'task_3', 'attempt_4')
         request_mock.assert_called_with(
-            '/proxy/app_1/ws/v1/mapreduce/jobs/job_2/tasks/task_3/attempt/attempt_4/counters')
+            '/proxy/app_1/ws/v1/mapreduce/jobs/job_2/tasks/task_3/attempts/attempt_4/counters'
+        )
diff --git a/yarn_api_client/application_master.py b/yarn_api_client/application_master.py
@@ -73,16 +73,21 @@ def job(self, application_id, job_id):
 
         return self.request(path)
 
-    def job_attempts(self, job_id):
+    def job_attempts(self, application_id, job_id):
         """
         With the job attempts API, you can obtain a collection of resources
         that represent the job attempts.
 
+        :param str application_id: The application id
         :param str job_id: The job id
         :returns: API response object with JSON data
         :rtype: :py:class:`yarn_api_client.base.Response`
         """
-        pass
+
+        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/jobattempts'.format(
+            appid=application_id, jobid=job_id)
+
+        return self.request(path)
 
     def job_counters(self, application_id, job_id):
         """
@@ -189,12 +194,49 @@ def task_attempt(self, application_id, job_id, task_id, attempt_id):
         :returns: API response object with JSON data
         :rtype: :py:class:`yarn_api_client.base.Response`
         """
-        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/tasks/{taskid}/attempt/{attemptid}'.format(
+        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/tasks/{taskid}/attempts/{attemptid}'.format(
             appid=application_id, jobid=job_id, taskid=task_id,
             attemptid=attempt_id)
 
         return self.request(path)
 
+    def task_attempt_state(self, application_id, job_id, task_id, attempt_id):
+        """
+        With the task attempt state API, you can query the state of a submitted
+        task attempt.
+
+        :param str application_id: The application id
+        :param str job_id: The job id
+        :param str task_id: The task id
+        :param str attempt_id: The attempt id
+        :returns: API response object with JSON data
+        :rtype: :py:class:`yarn_api_client.base.Response`
+        """
+        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/tasks/{taskid}/attempts/{attemptid}/state'.format(
+            appid=application_id, jobid=job_id, taskid=task_id,
+            attemptid=attempt_id)
+
+        return self.request(path)
+
+    def task_attempt_state_kill(self, application_id, job_id, task_id, attempt_id):
+        """
+        Kill specific attempt using task attempt state API.
+
+        :param str application_id: The application id
+        :param str job_id: The job id
+        :param str task_id: The task id
+        :param str attempt_id: The attempt id
+        :returns: API response object with JSON data
+        :rtype: :py:class:`yarn_api_client.base.Response`
+        """
+        data = {"state": "KILLED"}
+
+        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/tasks/{taskid}/attempts/{attemptid}/state'.format(
+            appid=application_id, jobid=job_id, taskid=task_id,
+            attemptid=attempt_id)
+
+        return self.request(path, 'PUT', data=data)
+
     def task_attempt_counters(self, application_id, job_id, task_id, attempt_id):
         """
         With the task attempt counters API, you can object a collection
@@ -207,7 +249,7 @@ def task_attempt_counters(self, application_id, job_id, task_id, attempt_id):
         :returns: API response object with JSON data
         :rtype: :py:class:`yarn_api_client.base.Response`
         """
-        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/tasks/{taskid}/attempt/{attemptid}/counters'.format(
+        path = '/proxy/{appid}/ws/v1/mapreduce/jobs/{jobid}/tasks/{taskid}/attempts/{attemptid}/counters'.format(
             appid=application_id, jobid=job_id, taskid=task_id,
             attemptid=attempt_id)