refactor

Signed-off-by: Javan Lacerda <javanlacerda@google.com>

Author: javanlacerda

src/clusterfuzz/_internal/k8s/job_template.yaml

@@ -15,16 +15,44 @@ apiVersion: batch/v1
 kind: Job
 metadata:
   name: "{{job_name}}"
+  labels:
+    task_name: "{{task_name}}"
+    job_name: "{{clusterfuzz_job_name}}"
 spec:
   ttlSecondsAfterFinished: 100
   activeDeadlineSeconds: {{active_deadline_seconds}}
   template:
+    metadata:
+      labels:
+        task_name: "{{task_name}}"
+        job_name: "{{clusterfuzz_job_name}}"
+        {% if is_kata %}
+        app.kubernetes.io/name: clusterfuzz-kata-job
+        {% endif %}
     spec:
+      {% if is_kata %}
+      runtimeClassName: kata
+      {% endif %}
+      dnsPolicy: ClusterFirstWithHostNet
       serviceAccountName: "{{service_account_name}}"
       containers:
       - name: "{{job_name}}"
         image: "{{docker_image}}"
         imagePullPolicy: IfNotPresent
+        lifecycle:
+          postStart:
+            exec:
+              command:
+              - /bin/sh
+              - -c
+              - mkdir -p /tmp/.X11-unix && chmod 1777 /tmp/.X11-unix
+        resources:
+          requests:
+            cpu: '2'
+            memory: 3.75Gi
+          limits:
+            cpu: '2'
+            memory: 3.75Gi
         env:
         - name: HOST_UID
           value: '1337'
@@ -57,5 +85,9 @@ spec:
         emptyDir:
           medium: Memory
           sizeLimit: 1.9Gi
+      {% if is_kata %}
+      nodeSelector:
+        cloud.google.com/gke-nodepool: kata-enabled-pool
+      {% endif %}
       restartPolicy: Never
-  backoffLimit: 0
+  backoffLimit: 0

src/clusterfuzz/_internal/k8s/kata_job_template.yaml

@@ -136,8 +136,7 @@ def _create_job_body(config: KubernetesJobConfig, input_url: str,
                      service_account_name: str) -> dict:
   """Creates the body of a Kubernetes job."""
 
-  job_name = config.job_type.replace('_', '-') + '-' + str(uuid.uuid4()).split(
-      '-', maxsplit=1)[0]
+  job_name = f'cf-job-{str(uuid.uuid4())}'
   job_name = job_name.lower()
 
   # Set up Jinja2 environment and load the template.
@@ -153,6 +152,9 @@ def _create_job_body(config: KubernetesJobConfig, input_url: str,
       'docker_image': config.docker_image,
       'clusterfuzz_release': config.clusterfuzz_release,
       'input_url': input_url,
+      'is_kata': config.is_kata,
+      'task_name': config.command,
+      'clusterfuzz_job_name': config.job_type,
   }
 
   # Render the template and load as YAML.
@@ -314,39 +316,6 @@ def create_utask_main_jobs(self, remote_tasks: typing.List[types.RemoteTask]):
     for config, input_urls in job_specs.items():
       # TODO(javanlacerda): Batch multiple tasks into a single job.
       for input_url in input_urls:
-        if config.is_kata:
-          jobs.append(self.create_kata_container_job(config, input_url))
-        else:
-          jobs.append(self.create_job(config, input_url))
+        jobs.append(self.create_job(config, input_url))
 
     return jobs
-
-  def create_kata_container_job(self, config: KubernetesJobConfig,
-                                input_url: str) -> str:
-    """Creates a Kubernetes job that runs in a Kata container."""
-    service_account_name = self._create_service_account_if_needed(
-        config.service_account_email)
-    job_name = 'clusterfuzz-kata-job-' + str(uuid.uuid4()).split(
-        '-', maxsplit=1)[0]
-
-    # Set up Jinja2 environment and load the template.
-    template_dir = os.path.dirname(__file__)
-    jinja_env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir))
-    template = jinja_env.get_template('kata_job_template.yaml')
-
-    # Define the context with all the dynamic values.
-    context = {
-        'job_name': job_name,
-        'active_deadline_seconds': tasks.get_task_duration(config.command),
-        'service_account_name': service_account_name,
-        'docker_image': config.docker_image,
-        'clusterfuzz_release': config.clusterfuzz_release,
-        'input_url': input_url,
-    }
-
-    # Render the template and load as YAML.
-    rendered_spec = template.render(context)
-    job_spec = yaml.safe_load(rendered_spec)
-
-    self._batch_api.create_namespaced_job(body=job_spec, namespace='default')
-    return job_name

src/clusterfuzz/_internal/k8s/service.py

@@ -220,40 +220,6 @@ def test_create_job(self, mock_get_logging_config_dict):
         namespace='default',
         body=k8s_client.V1DeleteOptions(propagation_policy='Foreground'))
 
-  @unittest.skip('Should be implemented against a cluster that supports kata')
-  def test_create_kata_container_job(self, mock_get_logging_config_dict):
-    """Tests creating a Kata container job."""
-    input_urls = []
-    actual_job_name = self.kubernetes_client.create_kata_container_job(
-        self.image, input_urls)
-
-    # Wait for the job to be created.
-    time.sleep(5)
-
-    job = self.api_client.read_namespaced_job(actual_job_name, 'default')
-    self.assertIsNotNone(job)
-    self.assertEqual(job.metadata.name, actual_job_name)
-    self.assertEqual(job.spec.template.spec.runtime_class_name, 'kata')
-
-    # Wait for the job to start running.
-    job_running = False
-    for _ in range(180):
-      job = self.api_client.read_namespaced_job(actual_job_name, 'default')
-      if job.status.active or job.status.succeeded:
-        job_running = True
-        break
-      time.sleep(1)
-
-    self.assertTrue(
-        job_running,
-        f'Kata Job {actual_job_name} did not start running. Status: {job.status}'
-    )
-
-    self.api_client.delete_namespaced_job(
-        name=actual_job_name,
-        namespace='default',
-        body=k8s_client.V1DeleteOptions(propagation_policy='Foreground'))
-
   @mock.patch('clusterfuzz._internal.k8s.service._get_k8s_job_configs')
   @mock.patch(
       'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module')

src/clusterfuzz/_internal/tests/core/k8s/k8s_service_e2e_test.py

@@ -52,7 +52,7 @@ def test_create_uworker_main_batch_jobs_limit_not_reached(
 
     # We expect this to proceed to job creation logic (which we mock to avoid actual creation)
     with mock.patch.object(service.KubernetesService,
-                           'create_kata_container_job') as mock_create:
+                           'create_job') as mock_create:
       kube_service.create_utask_main_jobs(
           [types.RemoteTask('fuzz', 'job1', 'url1')])
       self.assertTrue(mock_create.called)

src/clusterfuzz/_internal/tests/core/k8s/k8s_service_limit_test.py

@@ -48,10 +48,9 @@ def setUp(self):
         value=100.0).put()
 
   @mock.patch.object(service.KubernetesService, '_get_pending_jobs_count')
-  @mock.patch.object(service.KubernetesService, 'create_kata_container_job')
   @mock.patch.object(service.KubernetesService, 'create_job')
   def test_create_uworker_main_batch_jobs(
-      self, mock_create_job, mock_create_kata_job, mock_get_pending_count, _):
+      self, mock_create_job, mock_get_pending_count, _):
     """Tests the creation of uworker main batch jobs."""
     mock_get_pending_count.return_value = 0
     tasks = [
@@ -63,13 +62,11 @@ def test_create_uworker_main_batch_jobs(
     kube_service = service.KubernetesService()
     kube_service.create_utask_main_jobs(tasks)
 
-    # Assuming default config implies Kata, and no batching of URLs.
-    # Total 3 tasks, so 3 calls.
-    self.assertEqual(3, mock_create_kata_job.call_count)
-    self.assertEqual(0, mock_create_job.call_count)
+    # Total 3 tasks, so 3 calls to create_job.
+    self.assertEqual(3, mock_create_job.call_count)
 
     urls = sorted(
-        [call.args[1] for call in mock_create_kata_job.call_args_list])
+        [call.args[1] for call in mock_create_job.call_args_list])
     self.assertEqual(urls, ['url1', 'url2', 'url3'])
 
   @mock.patch('kubernetes.client.CoreV1Api')
@@ -106,8 +103,8 @@ def test_create_uworker_main_batch_jobs_limit_reached(
     self.assertTrue(mock_pubsub_task.do_not_ack)
 
   @mock.patch('kubernetes.client.BatchV1Api')
-  def test_create_kata_container_job_spec(self, mock_batch_api_cls, _):
-    """Tests that create_kata_container_job generates the correct spec."""
+  def test_create_job_kata(self, mock_batch_api_cls, _):
+    """Tests that create_job generates the correct spec for Kata."""
     mock_batch_api = mock_batch_api_cls.return_value
     kube_service = service.KubernetesService()
     # Mock _create_service_account_if_needed
@@ -123,7 +120,7 @@ def test_create_kata_container_job_spec(self, mock_batch_api_cls, _):
         clusterfuzz_release='prod',
         is_kata=True)
 
-    kube_service.create_kata_container_job(config, 'input_url')
+    kube_service.create_job(config, 'input_url')
 
     self.assertTrue(mock_batch_api.create_namespaced_job.called)
     call_args = mock_batch_api.create_namespaced_job.call_args
@@ -137,22 +134,18 @@ def test_create_kata_container_job_spec(self, mock_batch_api_cls, _):
     self.assertEqual(['ALL'],
                      container['securityContext']['capabilities']['add'])
 
-    # Check HOST_UID env var
-    env_names = {e['name']: e['value'] for e in container['env']}
-    self.assertEqual('1337', env_names['HOST_UID'])
-
-    # Check shm size
-    volumes = {v['name']: v for v in pod_spec['volumes']}
-    self.assertEqual('1.9Gi', volumes['dshm']['emptyDir']['sizeLimit'])
+    # Check labels
+    self.assertEqual('fuzz', job_body['metadata']['labels']['task_name'])
+    self.assertEqual('test-job', job_body['metadata']['labels']['job_name'])
 
-    # Check Service Account
-    self.assertEqual('test', pod_spec['serviceAccountName'])
-    kube_service._create_service_account_if_needed.assert_called_with(
-        'test@clusterfuzz-test.iam.gserviceaccount.com')
+    # Check Kata specific fields
+    self.assertEqual('kata', pod_spec['runtimeClassName'])
+    self.assertEqual('ClusterFirstWithHostNet', pod_spec['dnsPolicy'])
+    self.assertIn('lifecycle', container)
 
   @mock.patch('kubernetes.client.BatchV1Api')
-  def test_create_job(self, mock_batch_api_cls, _):
-    """Tests create_job."""
+  def test_create_job_standard(self, mock_batch_api_cls, _):
+    """Tests create_job for standard container."""
     mock_batch_api = mock_batch_api_cls.return_value
     kube_service = service.KubernetesService()
     kube_service._create_service_account_if_needed = mock.Mock(
@@ -173,11 +166,15 @@ def test_create_job(self, mock_batch_api_cls, _):
     call_args = mock_batch_api.create_namespaced_job.call_args
     job_body = call_args.kwargs['body']
 
+    pod_spec = job_body['spec']['template']['spec']
+    container = pod_spec['containers'][0]
+
     # Check Service Account
-    self.assertEqual('test-sa',
-                     job_body['spec']['template']['spec']['serviceAccountName'])
-    kube_service._create_service_account_if_needed.assert_called_with(
-        'test-email')
+    self.assertEqual('test-sa', pod_spec['serviceAccountName'])
+    
+    # Check that Kata specific fields are NOT present
+    self.assertNotIn('runtimeClassName', pod_spec)
+    self.assertIn('volumeMounts', container)
 
   @mock.patch(
       'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module')
@@ -194,4 +191,4 @@ def test_create_uworker_main_batch_job(self, mock_create_batch_jobs,
     self.assertEqual(1, len(tasks))
     self.assertEqual('command', tasks[0].command)
     self.assertEqual('job', tasks[0].job_type)
-    self.assertEqual('url', tasks[0].input_download_url)
+    self.assertEqual('url', tasks[0].input_download_url)

src/clusterfuzz/_internal/tests/core/k8s/k8s_service_test.py

@@ -77,7 +77,10 @@ def test_create_job(self):
     called_args, called_kwargs = self.k8s_client._batch_api.create_namespaced_job.call_args
     self.assertEqual(called_args, ())
     job_body = called_kwargs['body']
-    self.assertEqual(job_body['metadata']['name'], 'test-job-a0b1c2d3')
+    self.assertEqual(job_body['metadata']['name'],
+                     'cf-job-a0b1c2d3-e4f5-6789-0123-456789abcdef')
+    self.assertEqual(job_body['metadata']['labels']['task_name'], 'fuzz')
+    self.assertEqual(job_body['metadata']['labels']['job_name'], 'test-job')
     self.assertEqual(
         job_body['spec']['template']['spec']['containers'][0]['image'],
         'test-image')

src/clusterfuzz/_internal/tests/core/kubernetes/kubernetes_test.py

@@ -51,7 +51,7 @@
 _LICENSE_CHECK_IGNORE = 'LICENSE_CHECK_IGNORE'
 _PY_TEST_SUFFIX = '_test.py'
 _PY_INIT_FILENAME = '__init__.py'
-_YAML_EXCEPTIONS = ['bad.yaml', 'job_template.yaml', 'kata_job_template.yaml']
+_YAML_EXCEPTIONS = ['bad.yaml', 'job_template.yaml']
 
 _error_occurred = False