diff --git a/batch/spellcheck/README.md b/batch/spellcheck/README.md index 975f63733e..fb2898d4d4 100644 --- a/batch/spellcheck/README.md +++ b/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 diff --git a/robotoff/batch/buckets.py b/robotoff/batch/buckets.py index 19c5ae1d02..d27b1c2d73 100644 --- a/robotoff/batch/buckets.py +++ b/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 diff --git a/robotoff/batch/extraction.py b/robotoff/batch/extraction.py index d4f517836f..bc75be9d95 100644 --- a/robotoff/batch/extraction.py +++ b/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. diff --git a/robotoff/batch/launch.py b/robotoff/batch/launch.py index 8b00f097aa..d8467c6e0c 100644 --- a/robotoff/batch/launch.py +++ b/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.