Tweakable batch inference params (#21)

Author: seanshi-scale

.flake8

@@ -4,6 +4,8 @@ max-line-length = 79
 max-complexity = 18
 select = B,C,E,F,W,T4,B9
 exclude =
+ # clientlib is documentation only, not runnable code
+ launch/clientlib
  # All of these excludes should mirror something in .gitignore
  .git,
  __pychache__,

.pylintrc

@@ -24,3 +24,7 @@ reports=no
 
 [tool.pylint.FORMAT]
 max-line-length=79
+
+[MASTER]
+# Ignore anything inside launch/clientlib (since it's documentation)
+ignore=clientlib

launch/client.py

@@ -791,8 +791,9 @@ def batch_async_request(
  self,
  bundle_name: str,
  urls: List[str],
- batch_url_file_location: str = None,
+ batch_url_file_location: Optional[str] = None,
  serialization_format: str = "json",
+ batch_task_options: Optional[Dict[str, Any]] = None,
  ):
  """
  Sends a batch inference request to the Model Endpoint at endpoint_id, returns a key that can be used to retrieve
@@ -806,10 +807,33 @@ def batch_async_request(
  Must be accessible by Scale Launch, hence urls need to either be public or signedURLs.
  batch_url_file_location: In self-hosted mode, the input to the batch job will be uploaded
  to this location if provided. Otherwise, one will be determined from bundle_location_fn()
+ batch_task_options: A Dict of optional endpoint/batch task settings, i.e. certain endpoint settings
+ like cpus, memory, gpus, gpu_type, max_workers, as well as under-the-hood batch job settings, like
+ pyspark_partition_size, pyspark_max_executors.
 
  Returns:
  An id/key that can be used to fetch inference results at a later time
  """
+
+ if batch_task_options is None:
+ batch_task_options = {}
+ allowed_batch_task_options = {
+ "cpus",
+ "memory",
+ "gpus",
+ "gpu_type",
+ "max_workers",
+ "pyspark_partition_size",
+ "pyspark_max_executors",
+ }
+ if (
+ len(set(batch_task_options.keys()) - allowed_batch_task_options)
+ > 0
+ ):
+ raise ValueError(
+ f"Disallowed options {set(batch_task_options.keys()) - allowed_batch_task_options} for batch task"
+ )
+
  f = StringIO()
  make_batch_input_file(urls, f)
  f.seek(0)
@@ -837,6 +861,7 @@ def batch_async_request(
  input_path=file_location,
  serialization_format=serialization_format,
  )
+ payload.update(batch_task_options)
  payload = self.endpoint_auth_decorator_fn(payload)
  resp = self.connection.post(
  route=f"{BATCH_TASK_PATH}/{bundle_name}",

pyproject.toml

@@ -19,6 +19,11 @@ exclude = '''
 )
 '''
 
+[tool.mypy]
+exclude = [
+ '^launch/clientlib/'
+]
+
 [tool.poetry]
 name = "scale-launch"
 version = "0.1.0"