Update spec format to include arch (#120)

* update spec format to include arch

* add tests

* fix typo

docs/api.md

@@ -43,7 +43,7 @@ Given a spec, the API will calculate the optimal resource allocation for the job
 The spec sent to the endpoint should have the following format:
 
 ```
-pkg_name@pkg_version +variant1+variant2%compiler@compiler_version
+pkg_name@pkg_version +variant1+variant2 arch=arch%compiler@compiler_version
 ```
 
 Be sure that the string is URL-encoded. For instance, the `urllib.parse.quote` method will ensure the proper format. Without it, the allocation algorithm may return inaccurate results.

gantry/tests/defs/prediction.py

@@ -4,18 +4,18 @@
 from gantry.util.spec import parse_alloc_spec
 
 NORMAL_BUILD = parse_alloc_spec(
-    "[email protected] ~caffe2+cuda+cudnn~debug+distributed+fbgemm+gloo+kineto~magma~metal+mkldnn+mpi~nccl+nnpack+numa+numpy+onnx_ml+openmp+qnnpack~rocm+tensorpipe~test+valgrind+xnnpack build_system=python_pip cuda_arch=80%[email protected]"
+    "[email protected] ~caffe2+cuda+cudnn~debug+distributed+fbgemm+gloo+kineto~magma~metal+mkldnn+mpi~nccl+nnpack+numa+numpy+onnx_ml+openmp+qnnpack~rocm+tensorpipe~test+valgrind+xnnpack build_system=python_pip cuda_arch=80 arch=x86_64%[email protected]"
 )
 
 # everything in NORMAL_BUILD["package"]["variants"] except removing build_system=python_pip
 # in order to test the expensive variants filter
 EXPENSIVE_VARIANT_BUILD = parse_alloc_spec(
-    "[email protected] ~caffe2+cuda+cudnn~debug+distributed+fbgemm+gloo+kineto~magma~metal+mkldnn+mpi~nccl+nnpack+numa+numpy+onnx_ml+openmp+qnnpack~rocm+tensorpipe~test+valgrind+xnnpack cuda_arch=80%[email protected]"
+    "[email protected] ~caffe2+cuda+cudnn~debug+distributed+fbgemm+gloo+kineto~magma~metal+mkldnn+mpi~nccl+nnpack+numa+numpy+onnx_ml+openmp+qnnpack~rocm+tensorpipe~test+valgrind+xnnpack cuda_arch=80 arch=x86_64%[email protected]"
 )
 
 # no variants should match this, so we expect the default prediction
 BAD_VARIANT_BUILD = parse_alloc_spec(
-    "[email protected] +no~expensive~variants+match%[email protected]"
+    "[email protected] +no~expensive~variants+match arch=x86_64%[email protected]"
 )
 
 # calculated by running the baseline prediction algorithm on the sample data in gantry/tests/sql/insert_prediction.sql

gantry/tests/test_prediction.py

@@ -76,13 +76,16 @@ async def test_empty_sample(db_conn):
 # Test validate_payload
 def test_valid_spec():
     """Tests that a valid spec is parsed correctly."""
-    assert parse_alloc_spec("[email protected] +json+native+treesitter%[email protected]") == {
+    assert parse_alloc_spec(
+        "[email protected] +json+native+treesitter arch=x86_64%[email protected]"
+    ) == {
         "pkg_name": "emacs",
         "pkg_version": "29.2-test",
         "pkg_variants": '{"json": true, "native": true, "treesitter": true}',
         "pkg_variants_dict": {"json": True, "native": True, "treesitter": True},
         "compiler_name": "gcc",
         "compiler_version": "12.3.0",
+        "arch": "x86_64",
     }
 
 

gantry/util/spec.py

@@ -43,16 +43,17 @@ def parse_alloc_spec(spec: str) -> dict:
     - pkg_variants_dict: dict
     - compiler: str
     - compiler_version: str
+    - arch: str
 
     Returns an empty dict if the spec is invalid.
 
     This format is specifically used for the allocation API and is documented
     for the client.
     """
 
-    # example: [email protected] +json+native+treesitter%[email protected]
+    # example: [email protected] +json+native+treesitter arch=x86_64%[email protected]
     # this regex accommodates versions made up of any non-space characters
-    spec_pattern = re.compile(r"(.+?)@(\S+)\s+(.+?)%([\w-]+)@(\S+)")
+    spec_pattern = re.compile(r"(.+?)@(\S+)\s+(.+?)\s+arch=(\S+)%([\w-]+)@(\S+)")
 
     match = spec_pattern.match(spec)
     if not match:
@@ -64,6 +65,7 @@ def parse_alloc_spec(spec: str) -> dict:
         pkg_name,
         pkg_version,
         pkg_variants,
+        arch,
         compiler_name,
         compiler_version,
     ) = match.groups()
@@ -83,6 +85,7 @@ def parse_alloc_spec(spec: str) -> dict:
         "pkg_variants_dict": pkg_variants_dict,
         "compiler_name": compiler_name,
         "compiler_version": compiler_version,
+        "arch": arch,
     }
 
     return spec_dict

gantry/views.py

@@ -46,7 +46,7 @@ async def allocation(request: web.Request) -> web.Response:
     that set resource allocations based on historical data.
 
     acceptable spec format:
-    pkg_name@pkg_version +variant1+variant2%compiler@compiler_version
+    pkg_name@pkg_version +variant1+variant2%compiler arch=arch@compiler_version
     NOTE: there must be a space between the package version and the variants
 
     returns: