refactor(batch-spellcheck): 💚 Fix some bugs: batch-extraction & batch…

…-launch

batch/spellcheck/README.md

@@ -11,6 +11,7 @@
     * A100: a2-highgpu-1g: $3.748064
 * A100/Cuda doesn't support FP8
 * A100 has less availability than L4: need to wait for batch job (can be long)
+* Don't forget to enable **Batch & Storage API** if used without gcloud
 
 ## Links
 

robotoff/batch/buckets.py

@@ -35,9 +35,9 @@ def __init__(
         self.suffix_postprocess = suffix_postprocess
 
     @classmethod
-    def from_job_type(cls, job_type:BatchJobType) -> "GoogleStorageBucketForBatchJob":
-        """Initialize the class with the configuration file corresponding to the batch job type.
-        Useful to adapt bucket upload and download during the batch job process.
+    def from_job_type(cls, job_type: BatchJobType) -> "GoogleStorageBucketForBatchJob":
+        """Initialize the class with the bucket and suffix names corresponding to the batch job type.
+        Used to adapt bucket upload and download during the batch job process.
 
         :param job_type: Batch job type. 
         :type job_type: BatchJobType

robotoff/batch/extraction.py

@@ -31,7 +31,7 @@ def extract_from_dataset(
 
         :param job_type: Batch job type.
         :type job_type: BatchJobType
-        :param output_dir: Directory to save the extracted data.
+        :param output_dir: Directory to save the extracted data as a parquet file.
         :type output_dir: str
         :param dataset_path: Path to the jsonl.gz dataset.
         :type dataset_path: Path, optional. Default to settings.JSONL_DATASET_PATH. Mainly used for testing.

robotoff/batch/launch.py

@@ -3,9 +3,10 @@
 import enum
 import yaml
 import datetime
+import re
 
 from google.cloud import batch_v1
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict
 
 from robotoff import settings
 
@@ -14,7 +15,7 @@
 class BatchJobType(enum.Enum):
     """Each job type correspond to a task that will be executed in the batch job."""
 
-    ingredients_spellcheck = "ingredients_spellcheck"
+    ingredients_spellcheck = "ingredients-spellcheck"
 
 
 # Paths batch job config files
@@ -25,6 +26,8 @@ class BatchJobType(enum.Enum):
 
 class GoogleBatchJobConfig(BaseModel):
     """Batch job configuration class."""
+    # By default, extra fields are just ignored. We raise an error in case of extra fields.
+    model_config: ConfigDict = {"extra": "forbid"}
 
     job_name: str = Field(
         description="The name of the job. It needs to be unique amongst exisiting batch job names.",
@@ -33,6 +36,9 @@ class GoogleBatchJobConfig(BaseModel):
         pattern=r"^europe-west\d{1,2}$",
         description="The region in which the job will run. Regions that are available for Batch are listed on: https://cloud.google.com/compute/docs/gpus/gpu-regions-zones. We restrict to Europe-West for now.",
     )
+    container_image_uri: str = Field(
+        description="The URI of the container image to use for the job. SHould be a valid Image URI.",
+    )
     entrypoint: Optional[str] = Field(
         default=None,
         description="The entrypoint for the container. If None, use default entrypoint.",
@@ -100,12 +106,16 @@ def init(cls, job_type: BatchJobType):
         :param job_type: Batch job type.
         :type job_type: BatchJobType
         """
+        # Batch job name should respect a specific pattern, or returns an error
+        pattern = "^[a-z]([a-z0-9-]{0,61}[a-z0-9])?$"
+        if not re.match(pattern, job_type.value):
+            raise ValueError(f"Job name should respect the pattern: {pattern}. Current job name: {job_type.value}")
+
         # Generate unique id for the job
         unique_job_name = (
-            job_type.name + "-" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+            job_type.value + "-" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
         )
-
-        # Load config from job_type
+        # Load config file from job_type
         config_path = BATCH_JOB_TYPE_TO_CONFIG_PATH[job_type]
         with open(config_path, "r") as f:
             config = yaml.safe_load(f)
@@ -135,8 +145,10 @@ def launch_job(
     ) -> batch_v1.Job:
         """This method creates a Batch Job on GCP.
 
-        Method copied from https://github.com/GoogleCloudPlatform/python-docs-samples/tree/main/batch/create
-
+        Sources:
+        * https://github.com/GoogleCloudPlatform/python-docs-samples/tree/main/batch/create
+        * https://cloud.google.com/python/docs/reference/batch/latest/google.cloud.batch_v1.types
+        
         :param google_batch_launch_config: Config to run a job on Google Batch.
         :type google_batch_launch_config: GoogleBatchLaunchConfig
         :param batch_job_config: Config to run a specific job on Google Batch.