rename types

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

Author: javanlacerda

src/clusterfuzz/_internal/batch/service.py

@@ -35,7 +35,7 @@
 from clusterfuzz._internal.datastore import ndb_utils
 from clusterfuzz._internal.google_cloud_utils import credentials
 from clusterfuzz._internal.metrics import logs
-from clusterfuzz._internal.remote_task import types
+from clusterfuzz._internal.remote_task import remote_task_types
 from clusterfuzz._internal.system import environment
 
 # A named tuple that defines the execution environment for a batch workload.
@@ -170,7 +170,8 @@ def count_queued_or_scheduled_tasks(project: str,
   """Counts the number of queued and scheduled tasks."""
   region = f'projects/{project}/locations/{region}'
   jobs_filter = 'Status.State="SCHEDULED" OR Status.State="QUEUED"'
-  req = batch.types.ListJobsRequest(parent=region, filter=jobs_filter)
+  req = batch.remote_task_types.ListJobsRequest(
+      parent=region, filter=jobs_filter)
   queued = 0
   scheduled = 0
   for job in _batch_client().list_jobs(request=req):
@@ -193,13 +194,14 @@ def is_remote_task(command: str, job_name: str) -> bool:
   be found for the given command and job type.
   """
   try:
-    _get_specs_from_config([types.RemoteTask(command, job_name, None)])
+    _get_specs_from_config(
+        [remote_task_types.RemoteTask(command, job_name, None)])
     return True
   except ValueError:
     return False
 
 
-def _get_config_names(batch_tasks: List[types.RemoteTask]):
+def _get_config_names(batch_tasks: List[remote_task_types.RemoteTask]):
   """"Gets the name of the configs for each batch_task. Returns a dict
   that is indexed by command and job_type for efficient lookup."""
   job_names = {task.job_type for task in batch_tasks}
@@ -253,7 +255,8 @@ def _get_subconfig(batch_config, instance_spec):
   return all_subconfigs[weighted_subconfig.name]
 
 
-def _get_specs_from_config(batch_tasks: List[types.RemoteTask]) -> Dict:
+def _get_specs_from_config(
+    batch_tasks: List[remote_task_types.RemoteTask]) -> Dict:
   """Gets the configured specifications for a batch workload."""
   if not batch_tasks:
     return {}
@@ -321,7 +324,7 @@ def _get_specs_from_config(batch_tasks: List[types.RemoteTask]) -> Dict:
   return specs
 
 
-class GcpBatchService(types.RemoteTaskInterface):
+class GcpBatchService(remote_task_types.RemoteTaskInterface):
   """A high-level service for creating and managing remote tasks.
   
   This service provides a simple interface for scheduling ClusterFuzz tasks on
@@ -371,13 +374,16 @@ def create_utask_main_job(self, module: str, job_type: str,
                             input_download_url: str):
     """Creates a single batch job for a uworker main task."""
     command = task_utils.get_command_from_module(module)
-    batch_tasks = [types.RemoteTask(command, job_type, input_download_url)]
+    batch_tasks = [
+        remote_task_types.RemoteTask(command, job_type, input_download_url)
+    ]
     result = self.create_utask_main_jobs(batch_tasks)
     if result is None:
       return result
     return result[0]
 
-  def create_utask_main_jobs(self, remote_tasks: List[types.RemoteTask]):
+  def create_utask_main_jobs(self,
+                             remote_tasks: List[remote_task_types.RemoteTask]):
     """Creates a batch job for a list of uworker main tasks.
     
     This method groups the tasks by their workload specification and creates a

src/clusterfuzz/_internal/k8s/service.py

@@ -34,7 +34,7 @@
 from clusterfuzz._internal.datastore import feature_flags
 from clusterfuzz._internal.datastore import ndb_utils
 from clusterfuzz._internal.metrics import logs
-from clusterfuzz._internal.remote_task import types
+from clusterfuzz._internal.remote_task import remote_task_types
 from clusterfuzz._internal.system import environment
 
 CLUSTER_NAME = 'clusterfuzz-cronjobs-gke'
@@ -50,7 +50,7 @@
 ])
 
 
-def _get_config_names(remote_tasks: typing.List[types.RemoteTask]):
+def _get_config_names(remote_tasks: typing.List[remote_task_types.RemoteTask]):
   """"Gets the name of the configs for each batch_task. Returns a dict
 
   that is indexed by command and job_type for efficient lookup."""
@@ -62,7 +62,7 @@ def _get_config_names(remote_tasks: typing.List[types.RemoteTask]):
   config_map = {}
   for task in remote_tasks:
     if task.job_type not in job_map:
-      logs.error(f'{task.job_type} doesn\'t exist.')
+      logs.error(f"{task.job_type} doesn't exist.")
       continue
     if task.command == 'fuzz':
       suffix = '-PREEMPTIBLE-UNPRIVILEGED'
@@ -88,13 +88,13 @@ def _get_config_names(remote_tasks: typing.List[types.RemoteTask]):
 
 
 def _get_k8s_job_configs(
-    remote_tasks: typing.List[types.RemoteTask]
+    remote_tasks: typing.List[remote_task_types.RemoteTask]
 ) -> typing.Dict[typing.Tuple[str, str], KubernetesJobConfig]:
   """Gets the configured specifications for a batch workload."""
 
   if not remote_tasks:
     return {}
-  #TODO(javanlacerda): Create remote task config
+  # TODO(javanlacerda): Create remote task config
   batch_config = local_config.BatchConfig()
   config_map = _get_config_names(remote_tasks)
   configs = {}
@@ -162,7 +162,7 @@ def _create_job_body(config: KubernetesJobConfig, input_url: str,
   return yaml.safe_load(rendered_spec)
 
 
-class KubernetesService(types.RemoteTaskInterface):
+class KubernetesService(remote_task_types.RemoteTaskInterface):
   """A remote task execution client for Kubernetes."""
 
   def __init__(self, k8s_config_loaded: bool = False):
@@ -178,7 +178,7 @@ def _load_gke_credentials(self):
     credentials, _ = google.auth.default()
     project = utils.get_application_id()
     service = discovery.build('container', 'v1', credentials=credentials)
-    parent = f'projects/{project}/locations/-'
+    parent = f"projects/{project}/locations/-"
 
     try:
       # pylint: disable=no-member
@@ -188,12 +188,12 @@ def _load_gke_credentials(self):
       cluster = next((c for c in clusters if c['name'] == CLUSTER_NAME), None)
 
       if not cluster:
-        logs.error(f'Cluster {CLUSTER_NAME} not found in project {project}.')
-        print(f'DEBUG: Cluster {CLUSTER_NAME} not found in project {project}.')
+        logs.error(f"Cluster {CLUSTER_NAME} not found in project {project}.")
+        print(f"DEBUG: Cluster {CLUSTER_NAME} not found in project {project}.")
         return
 
     except Exception as e:
-      logs.error(f'Failed to list clusters in {project}: {e}')
+      logs.error(f"Failed to list clusters in {project}: {e}")
       return
 
     endpoint = cluster['endpoint']
@@ -214,13 +214,13 @@ def get_token(creds):
       request = google_requests.Request()
       if not creds.valid or creds.expired:
         creds.refresh(request)
-      return {'authorization': 'Bearer ' + creds.token}
+      return {"authorization": "Bearer " + creds.token}
 
     configuration.refresh_api_key_hook = lambda _: get_token(credentials)
     configuration.api_key = get_token(credentials)
 
     k8s_client.Configuration.set_default(configuration)
-    logs.info('GKE credentials loaded successfully.')
+    logs.info("GKE credentials loaded successfully.")
 
   def _create_service_account_if_needed(self,
                                         service_account_email: str) -> str:
@@ -266,36 +266,39 @@ def _get_pending_jobs_count(self) -> int:
           namespace='default',
           label_selector='app.kubernetes.io/name=clusterfuzz-kata-job',
           field_selector='status.phase=Pending')
-      logs.info(f'Found {len(pods.items)} pending jobs.')
+      logs.info(f"Found {len(pods.items)} pending jobs.")
       return len(pods.items)
     except Exception as e:
-      logs.error(f'Failed to list pods: {e}')
+      logs.error(f"Failed to list pods: {e}")
       return 0
 
   def create_utask_main_job(self, module: str, job_type: str,
                             input_download_url: str):
     """Creates a single batch job for a uworker main task."""
 
     command = task_utils.get_command_from_module(module)
-    batch_tasks = [types.RemoteTask(command, job_type, input_download_url)]
+    batch_tasks = [
+        remote_task_types.RemoteTask(command, job_type, input_download_url)
+    ]
     result = self.create_utask_main_jobs(batch_tasks)
 
     if result is None:
       return result
     return result[0]
 
-  def create_utask_main_jobs(self, remote_tasks: typing.List[types.RemoteTask]):
+  def create_utask_main_jobs(
+      self, remote_tasks: typing.List[remote_task_types.RemoteTask]):
     """Creates a batch job for a list of uworker main tasks.
 
     This method groups the tasks by their workload specification and creates a
     separate batch job for each group. This allows tasks with similar
     requirements to be processed together, which can improve efficiency.
     """
-    if feature_flags.FeatureFlags.K8S_PENDING_JOBS_LIMITER.enabled and \
-    feature_flags.FeatureFlags.K8S_PENDING_JOBS_LIMITER.content is not None\
-    and self._get_pending_jobs_count() >= int(feature_flags
-                                              .FeatureFlags.
-                                              K8S_PENDING_JOBS_LIMITER.content):
+    limit = None
+    if feature_flags.FeatureFlags.K8S_PENDING_JOBS_LIMITER.enabled:
+      limit = feature_flags.FeatureFlags.K8S_PENDING_JOBS_LIMITER.content
+
+    if limit is not None and self._get_pending_jobs_count() >= int(limit):
       logs.warning(
           f'Kubernetes job limit reached. Not acking {len(remote_tasks)} tasks.'
       )
@@ -310,11 +313,10 @@ def create_utask_main_jobs(self, remote_tasks: typing.List[types.RemoteTask]):
       logs.info(f'Scheduling {remote_task.command}, {remote_task.job_type}.')
       config = configs[(remote_task.command, remote_task.job_type)]
       job_specs[config].append(remote_task.input_download_url)
-    logs.info('Creating Kubernetes jobs.')
+    logs.info('Creating batch jobs.')
     jobs = []
-    logs.info('Batching k8s utask_mains.')
+    logs.info('Batching utask_mains.')
     for config, input_urls in job_specs.items():
-      # TODO(javanlacerda): Batch multiple tasks into a single job.
       for input_url in input_urls:
         jobs.append(self.create_job(config, input_url))
 

src/clusterfuzz/_internal/remote_task/__init__.py

@@ -26,10 +26,10 @@
 
 from clusterfuzz._internal.datastore import feature_flags
 from clusterfuzz._internal.metrics import logs
-from clusterfuzz._internal.remote_task import types
+from clusterfuzz._internal.remote_task import remote_task_types
 
 
-class RemoteTaskGate(types.RemoteTaskInterface):
+class RemoteTaskGate(remote_task_types.RemoteTaskInterface):
   """A generic dispatcher for remote task execution.
 
   This class acts as a high-level manager that abstracts away the specific
@@ -116,7 +116,8 @@ def create_utask_main_job(self, module: str, job_type: str,
     service = self._service_map[adapter_id]
     return service.create_utask_main_job(module, job_type, input_download_url)
 
-  def create_utask_main_jobs(self, remote_tasks: List[types.RemoteTask]):
+  def create_utask_main_jobs(self,
+                             remote_tasks: List[remote_task_types.RemoteTask]):
     """Creates a batch of remote tasks, distributing them across backends.
 
     This method handles two cases:

…usterfuzz/_internal/remote_task/types.py …nternal/remote_task/remote_task_types.pysrc/clusterfuzz/_internal/remote_task/types.py renamed to src/clusterfuzz/_internal/remote_task/remote_task_types.py src/clusterfuzz/_internal/remote_task/types.py renamed to src/clusterfuzz/_internal/remote_task/remote_task_types.py

@@ -22,7 +22,7 @@
 
 from clusterfuzz._internal.batch import service as batch_service
 from clusterfuzz._internal.datastore import data_types
-from clusterfuzz._internal.remote_task import types
+from clusterfuzz._internal.remote_task import remote_task_types
 from clusterfuzz._internal.tests.test_libs import helpers
 from clusterfuzz._internal.tests.test_libs import test_utils
 
@@ -199,9 +199,9 @@ def test_create_uworker_main_batch_jobs(self):
           ('command2', 'job2'): spec2,
       }
       tasks = [
-          types.RemoteTask('command1', 'job1', 'url1'),
-          types.RemoteTask('command1', 'job1', 'url2'),
-          types.RemoteTask('command2', 'job2', 'url3'),
+          remote_task_types.RemoteTask('command1', 'job1', 'url1'),
+          remote_task_types.RemoteTask('command1', 'job1', 'url2'),
+          remote_task_types.RemoteTask('command2', 'job2', 'url3'),
       ]
 
       # Call the function.
@@ -374,14 +374,14 @@ def test_get_specs_from_config_disk_size(self):
         name='libfuzzer_asan_test').put()
 
     spec = batch_service._get_specs_from_config(
-        [types.RemoteTask('fuzz', 'libfuzzer_asan_test', None)])
+        [remote_task_types.RemoteTask('fuzz', 'libfuzzer_asan_test', None)])
     self.assertEqual(spec['fuzz', 'libfuzzer_asan_test'].disk_size_gb, size)
 
   def test_get_specs_from_config_no_disk_size(self):
     """Test that disk_size_gb isn't mandatory."""
     data_types.Job(platform='LINUX', name='libfuzzer_asan_test').put()
     spec = batch_service._get_specs_from_config(
-        [types.RemoteTask('fuzz', 'libfuzzer_asan_test', None)])
+        [remote_task_types.RemoteTask('fuzz', 'libfuzzer_asan_test', None)])
     conf = batch_service._get_batch_config()
     expected_size = (
         conf.get('mapping')['LINUX-PREEMPTIBLE-UNPRIVILEGED']['disk_size_gb'])
@@ -406,7 +406,7 @@ def test_get_specs_from_config_with_disk_size_override(self):
         name=job_name).put()
 
     spec = batch_service._get_specs_from_config(
-        [types.RemoteTask('fuzz', job_name, None)])
+        [remote_task_types.RemoteTask('fuzz', job_name, None)])
     self.assertEqual(spec['fuzz', job_name].disk_size_gb, overridden_size)
 
   @mock.patch('clusterfuzz._internal.batch.service.utils.is_oss_fuzz')
@@ -422,7 +422,7 @@ def test_get_config_names_os_version(self, mock_get_all_from_query,
         name='job1', platform='LINUX', base_os_version='job-os-ubuntu-20')
     mock_get_all_from_query.return_value = [job1]
     config_map = batch_service._get_config_names(
-        [types.RemoteTask('fuzz', 'job1', None)])
+        [remote_task_types.RemoteTask('fuzz', 'job1', None)])
     self.assertEqual(config_map[('fuzz', 'job1')][2], 'job-os-ubuntu-20')
 
     # Test Case 2: OSS-Fuzz project, project-level version overrides job-level.
@@ -437,23 +437,23 @@ def test_get_config_names_os_version(self, mock_get_all_from_query,
     mock_get_all_from_query.return_value = [job2]
     mock_oss_fuzz_project_query.return_value.get.return_value = project
     config_map = batch_service._get_config_names(
-        [types.RemoteTask('fuzz', 'job2', None)])
+        [remote_task_types.RemoteTask('fuzz', 'job2', None)])
     self.assertEqual(config_map[('fuzz', 'job2')][2], 'project-os-ubuntu-24')
 
     # Test Case 3: OSS-Fuzz project, only project-level version exists.
     job3 = data_types.Job(name='job3', project='my-project', platform='LINUX')
     mock_get_all_from_query.return_value = [job3]
     mock_oss_fuzz_project_query.return_value.get.return_value = project
     config_map = batch_service._get_config_names(
-        [types.RemoteTask('fuzz', 'job3', None)])
+        [remote_task_types.RemoteTask('fuzz', 'job3', None)])
     self.assertEqual(config_map[('fuzz', 'job3')][2], 'project-os-ubuntu-24')
 
     # Test Case 4: Internal project, no version is set, should be None.
     mock_is_oss_fuzz.return_value = False
     job4 = data_types.Job(name='job4', platform='LINUX')
     mock_get_all_from_query.return_value = [job4]
     config_map = batch_service._get_config_names(
-        [types.RemoteTask('fuzz', 'job4', None)])
+        [remote_task_types.RemoteTask('fuzz', 'job4', None)])
     self.assertIsNone(config_map[('fuzz', 'job4')][2])
 
     # Test Case 5: OSS-Fuzz project, but no versions are set anywhere.
@@ -464,11 +464,11 @@ def test_get_config_names_os_version(self, mock_get_all_from_query,
     mock_get_all_from_query.return_value = [job5]
     mock_oss_fuzz_project_query.return_value.get.return_value = project_no_version
     config_map = batch_service._get_config_names(
-        [types.RemoteTask('fuzz', 'job5', None)])
+        [remote_task_types.RemoteTask('fuzz', 'job5', None)])
     self.assertIsNone(config_map[('fuzz', 'job5')][2])
 
 
 def _get_spec_from_config(command, job_name):
   return list(
       batch_service._get_specs_from_config(
-          [types.RemoteTask(command, job_name, None)]).values())[0]
+          [remote_task_types.RemoteTask(command, job_name, None)]).values())[0]

src/clusterfuzz/_internal/tests/core/batch/batch_service_test.py

@@ -29,7 +29,7 @@
 from clusterfuzz._internal.datastore import data_types
 from clusterfuzz._internal.k8s import service as kubernetes_service
 from clusterfuzz._internal.k8s.service import KubernetesJobConfig
-from clusterfuzz._internal.remote_task import types
+from clusterfuzz._internal.remote_task import remote_task_types
 from clusterfuzz._internal.tests.test_libs import test_utils
 
 
@@ -182,7 +182,7 @@ def tearDownClass(cls):
   def test_create_job(self, mock_get_logging_config_dict):
     """Tests creating a job."""
     input_url = 'url'
-    task = types.RemoteTask(None, 'test-job', None)
+    task = remote_task_types.RemoteTask(None, 'test-job', None)
     task.docker_image = self.image
 
     config = KubernetesJobConfig(
@@ -298,8 +298,8 @@ def test_create_uworker_main_batch_jobs(self, mock_get_command_from_module,
     }
 
     tasks = [
-        types.RemoteTask('fuzz', 'test-job1', 'url1'),
-        types.RemoteTask('fuzz', 'test-job2', 'url2'),
+        remote_task_types.RemoteTask('fuzz', 'test-job1', 'url1'),
+        remote_task_types.RemoteTask('fuzz', 'test-job2', 'url2'),
     ]
 
     actual_job_names = \