Wrappers around autogen client for docker image batch jobs (#93)

* things to fill in * untested bundle stuff * batch job stuff + docs, also untested * aaa * wrong function mb * fix some bugs * mypy pls * comment the jank

Author: seanshi-scale

launch/client.py

@@ -31,6 +31,12 @@
 from launch.api_client.model.create_batch_job_v1_request import (
  CreateBatchJobV1Request,
 )
+from launch.api_client.model.create_docker_image_batch_job_bundle_v1_request import (
+ CreateDockerImageBatchJobBundleV1Request,
+)
+from launch.api_client.model.create_docker_image_batch_job_v1_request import (
+ CreateDockerImageBatchJobV1Request,
+)
 from launch.api_client.model.create_model_bundle_v1_request import (
  CreateModelBundleV1Request,
 )
@@ -61,6 +67,9 @@
 from launch.api_client.model.triton_enhanced_runnable_image_flavor import (
  TritonEnhancedRunnableImageFlavor,
 )
+from launch.api_client.model.update_docker_image_batch_job_v1_request import (
+ UpdateDockerImageBatchJobV1Request,
+)
 from launch.api_client.model.update_model_endpoint_v1_request import (
  UpdateModelEndpointV1Request,
 )
@@ -73,6 +82,11 @@
  SCALE_LAUNCH_ENDPOINT,
  SCALE_LAUNCH_V1_ENDPOINT,
 )
+from launch.docker_image_batch_job_bundle import (
+ CreateDockerImageBatchJobBundleResponse,
+ DockerImageBatchJobBundleResponse,
+ ListDockerImageBatchJobBundleResponse,
+)
 from launch.find_packages import find_packages_from_imports, get_imports
 from launch.hooks import PostInferenceHooks
 from launch.make_batch_file import (
@@ -2074,6 +2088,255 @@ def get_batch_async_response(self, batch_job_id: str) -> Dict[str, Any]:
  resp = json.loads(response.response.data)
  return resp
 
+ def create_docker_image_batch_job_bundle(
+ self,
+ *,
+ name: str,
+ image_repository: str,
+ image_tag: str,
+ command: List[str],
+ env: Optional[Dict[str, str]],
+ mount_location: Optional[str],
+ cpus: Optional[int] = None,
+ memory: Optional[str] = None,
+ gpus: Optional[int] = None,
+ gpu_type: Optional[str] = None,
+ storage: Optional[str] = None,
+ ) -> CreateDockerImageBatchJobBundleResponse:
+ """
+ For self hosted mode only.
+
+ Creates a Docker Image Batch Job Bundle.
+
+ Parameters:
+ name:
+ A user-defined name for the bundle. Does not need to be unique.
+ image_repository:
+ The (short) repository of your image. For example, if your image is located at
+ 123456789012.dkr.ecr.us-west-2.amazonaws.com/repo:tag, and your version of Launch
+ is configured to look at 123456789012.dkr.ecr.us-west-2.amazonaws.com for Docker Images,
+ you would pass the value `repo` for the `image_repository` parameter.
+ image_tag:
+ The tag of your image inside of the repo. In the example above, you would pass
+ the value `tag` for the `image_tag` parameter.
+ command:
+ The command to run inside the docker image.
+ env:
+ A dictionary of environment variables to inject into your docker image.
+ mount_location:
+ A location in the filesystem where you would like a json-formatted file, controllable
+ on runtime, to be mounted. This allows behavior to be specified on runtime.
+ (Specifically, the contents of this file can be read via `json.load()` inside of the
+ user-defined code.)
+ cpus:
+ Optional default value for the number of cpus to give the job.
+ memory:
+ Optional default value for the amount of memory to give the job.
+ gpus:
+ Optional default value for the number of gpus to give the job.
+ gpu_type:
+ Optional default value for the type of gpu to give the job.
+ storage:
+ Optional default value for the amount of disk to give the job.
+ """
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ # Raw dictionary since it's not the entire request and we can get away with this
+ # Also having it be a CreateDockerImageBatchJobResourceRequest runs into problems
+ # if no values are specified
+ resource_requests = dict_not_none(
+ cpus=cpus,
+ memory=memory,
+ gpus=gpus,
+ gpu_type=gpu_type,
+ storage=storage,
+ )
+ create_docker_image_batch_job_bundle_request = CreateDockerImageBatchJobBundleV1Request(
+ **dict_not_none(
+ name=name,
+ image_repository=image_repository,
+ image_tag=image_tag,
+ command=command,
+ env=env,
+ mount_location=mount_location,
+ resource_requests=resource_requests,
+ )
+ )
+ response = api_instance.create_docker_image_batch_job_bundle_v1_docker_image_batch_job_bundles_post(
+ body=create_docker_image_batch_job_bundle_request, skip_deserialization=True
+ )
+ resp = CreateDockerImageBatchJobBundleResponse.parse_raw(response.response.data)
+ return resp
+
+ def get_docker_image_batch_job_bundle(
+ self, docker_image_batch_job_bundle_id: str
+ ) -> DockerImageBatchJobBundleResponse:
+ """
+ For self hosted mode only. Gets information for a single batch job bundle with a given id.
+ """
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ path_params = frozendict({"docker_image_batch_job_bundle_id": docker_image_batch_job_bundle_id})
+ response = api_instance.get_docker_image_batch_job_model_bundle_v1_docker_image_batch_job_bundles_docker_image_batch_job_bundle_id_get( # type: ignore # noqa: E501
+ path_params=path_params,
+ skip_deserialization=True,
+ )
+ resp = DockerImageBatchJobBundleResponse.parse_raw(response.response.data)
+
+ return resp
+
+ def get_latest_docker_image_batch_job_bundle(self, bundle_name: str) -> DockerImageBatchJobBundleResponse:
+ """
+ For self hosted mode only. Gets information for the latest batch job bundle with a given name.
+ """
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ query_params = frozendict({"bundle_name": bundle_name})
+ response = api_instance.get_latest_docker_image_batch_job_bundle_v1_docker_image_batch_job_bundles_latest_get( # type: ignore # noqa: E501
+ query_params=query_params,
+ skip_deserialization=True,
+ )
+ resp = DockerImageBatchJobBundleResponse.parse_raw(response.response.data)
+
+ return resp
+
+ def list_docker_image_batch_job_bundles(
+ self, bundle_name: Optional[str] = None, order_by: Optional[Literal["newest", "oldest"]] = None
+ ) -> ListDockerImageBatchJobBundleResponse:
+ """
+ For self hosted mode only. Gets information for multiple bundles.
+
+ Parameters:
+ bundle_name: The name of the bundles to retrieve. If not specified, this will retrieve all
+ bundles.
+ order_by: Either "newest", "oldest", or not specified. Specify to sort by newest/oldest.
+ """
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ query_params = frozendict(dict_not_none(bundle_name=bundle_name, order_by=order_by))
+ response = api_instance.list_docker_image_batch_job_model_bundles_v1_docker_image_batch_job_bundles_get( # type: ignore # noqa: E501
+ query_params=query_params,
+ skip_deserialization=True,
+ )
+ resp = ListDockerImageBatchJobBundleResponse.parse_raw(response.response.data)
+
+ return resp
+
+ def create_docker_image_batch_job(
+ self,
+ *,
+ labels: Dict[str, str],
+ docker_image_batch_job_bundle: Optional[Union[str, DockerImageBatchJobBundleResponse]] = None,
+ docker_image_batch_job_bundle_name: Optional[str] = None,
+ job_config: Optional[Dict[str, Any]] = None,
+ cpus: Optional[int] = None,
+ memory: Optional[str] = None,
+ gpus: Optional[int] = None,
+ gpu_type: Optional[str] = None,
+ storage: Optional[str] = None,
+ ):
+ """
+ For self hosted mode only.
+ Parameters:
+ docker_image_batch_job_bundle: Specifies the docker image bundle to use for the batch job.
+ Either the string id of a docker image bundle, or a
+ DockerImageBatchJobBundleResponse object.
+ Only one of docker_image_batch_job_bundle and docker_image_batch_job_bundle_name
+ can be specified.
+ docker_image_batch_job_bundle_name: The name of a batch job bundle. If specified,
+ Launch will use the most recent bundle with that name owned by the current user.
+ Only one of docker_image_batch_job_bundle and docker_image_batch_job_bundle_name
+ can be specified.
+ labels: Kubernetes labels that are present on the batch job.
+ job_config: A JSON-serializable python object that will get passed to the batch job,
+ specifically as the contents of a file mounted at `mount_location` inside the bundle.
+ You can call python's `json.load()` on the file to retrieve the contents.
+ cpus: Optional override for the number of cpus to give to your job. Either the default
+ must be specified in the bundle, or this must be specified.
+ memory: Optional override for the amount of memory to give to your job. Either the default
+ must be specified in the bundle, or this must be specified.
+ gpus: Optional number of gpus to give to the bundle. If not specified in the bundle or
+ here, will be interpreted as 0 gpus.
+ gpu_type: Optional type of gpu. If the final number of gpus is positive, must be specified
+ either in the bundle or here.
+ storage: Optional reserved amount of disk to give to your batch job. If not specified,
+ your job may be evicted if it is using too much disk.
+ """
+
+ assert (docker_image_batch_job_bundle is None) ^ (
+ docker_image_batch_job_bundle_name is None
+ ), "Exactly one of docker_image_batch_job_bundle and docker_image_batch_job_bundle_name must be specified"
+
+ if docker_image_batch_job_bundle is not None and isinstance(
+ docker_image_batch_job_bundle, DockerImageBatchJobBundleResponse
+ ):
+ docker_image_batch_job_bundle_id: Optional[str] = docker_image_batch_job_bundle.id
+ else:
+ docker_image_batch_job_bundle_id = docker_image_batch_job_bundle
+
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ # Raw dictionary since it's not the entire request and we can get away with this
+ # Also having it be a CreateDockerImageBatchJobResourceRequest runs into problems
+ # if no values are specified
+ resource_requests = dict_not_none(
+ cpus=cpus,
+ memory=memory,
+ gpus=gpus,
+ gpu_type=gpu_type,
+ storage=storage,
+ )
+ create_docker_image_batch_job_request = CreateDockerImageBatchJobV1Request(
+ **dict_not_none(
+ docker_image_batch_job_bundle_id=docker_image_batch_job_bundle_id,
+ docker_image_batch_job_bundle_name=docker_image_batch_job_bundle_name,
+ job_config=job_config,
+ labels=labels,
+ resource_requests=resource_requests,
+ )
+ )
+ response = api_instance.create_docker_image_batch_job_v1_docker_image_batch_jobs_post(
+ body=create_docker_image_batch_job_request,
+ skip_deserialization=True,
+ )
+ resp = json.loads(response.response.data)
+ return resp
+
+ def get_docker_image_batch_job(self, batch_job_id: str):
+ """
+ For self hosted mode only. Gets information about a batch job given a batch job id.
+ """
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ path_params = frozendict({"batch_job_id": batch_job_id})
+ response = (
+ api_instance.get_docker_image_batch_job_v1_docker_image_batch_jobs_batch_job_id_get( # type: ignore
+ path_params=path_params,
+ skip_deserialization=True,
+ )
+ )
+ resp = json.loads(response.response.data)
+
+ return resp
+
+ def update_docker_image_batch_job(self, batch_job_id: str, cancel: bool):
+ """
+ For self hosted mode only. Updates a batch job by id.
+ Use this if you want to cancel/delete a batch job.
+ """
+ with ApiClient(self.configuration) as api_client:
+ api_instance = DefaultApi(api_client)
+ path_params = frozendict({"batch_job_id": batch_job_id})
+ body = UpdateDockerImageBatchJobV1Request(cancel=cancel)
+ response = api_instance.update_docker_image_batch_job_v1_docker_image_batch_jobs_batch_job_id_put( # type: ignore # noqa: E501
+ body=body,
+ path_params=path_params,
+ skip_deserialization=True,
+ )
+ resp = json.loads(response.response.data)
+
+ return resp
+
 
 def _zip_directory(zipf: ZipFile, path: str) -> None:
  for root, _, files in os.walk(path):

launch/docker_image_batch_job_bundle.py

@@ -0,0 +1,60 @@
+import datetime
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+class CreateDockerImageBatchJobBundleResponse(BaseModel):
+ """
+ Response Object for creating a Docker Image Batch Job Bundle
+ Note: only available for self-hosted mode
+ """
+
+ docker_image_batch_job_bundle_id: str
+ """ID of the Docker Image Batch Job Bundle"""
+
+
+class DockerImageBatchJobBundleResponse(BaseModel):
+ """
+ Response object for a single Docker Image Batch Job Bundle
+ Note: only available for self-hosted mode
+ """
+
+ id: str
+ """ID of the Docker Image Batch Job Bundle"""
+ name: str
+ """Name of the Docker Image Batch Job Bundle"""
+ created_at: datetime.datetime
+ """Timestamp of when the Docker Image Batch Job Bundle was created"""
+ image_repository: str
+ """Short repository name of the underlying docker image"""
+ image_tag: str
+ """Tag of the underlying docker image"""
+ command: List[str]
+ """The command to run inside the docker image"""
+ env: Dict[str, str]
+ """Environment variables to be injected into the docker image"""
+ mount_location: Optional[str]
+ """Location of a json-formatted file to mount inside the docker image.
+ Contents get populated at runtime, and this is the method to change behavior on runtime."""
+ cpus: Optional[str]
+ """Default number of cpus to give to the docker image"""
+ memory: Optional[str]
+ """Default amount of memory to give to the docker image"""
+ storage: Optional[str]
+ """Default amount of disk to give to the docker image"""
+ gpus: Optional[int]
+ """Default number of gpus to give to the docker image"""
+ gpu_type: Optional[str]
+ """Default type of gpu, e.g. nvidia-tesla-t4, nvidia-ampere-a10 to give to the docker image"""
+
+
+class ListDockerImageBatchJobBundleResponse(BaseModel):
+ """
+ Response object for listing Docker Image Batch Job Bundles.
+ Note: only available for self-hosted mode
+ """
+
+ docker_image_batch_job_bundles: List[DockerImageBatchJobBundleResponse]
+ """A list of
+ [Docker Image Batch Job Bundles](./#launch.docker_image_batch_job_bundle.DockerImageBatchJobBundleResponse)."""