Planet-scale, fully managed artificial intelligence infrastructure service

A computerized method for managing artificial intelligence (AI) workloads in a cloud infrastructure platform, the computerized method comprising:

• receiving a first AI workload from a first tenant;
• scheduling the first AI workload for execution on a first resource selected from a group of heterogeneous distributed infrastructure resources, the first resource comprising first AI accelerator hardware;
• executing the first AI workload on the first resource based on the scheduling, whereby the first AI workload executes using the first AI accelerator hardware;
• executing a graphics processing unit (GPU)-related application programming interface (API) call that generates a physical handle to a GPU object, the GPU object comprising at least one of a stream, an event, or a library context;
• storing the GPU-related API call in a log;
• intercepting, by a checkpointing system, the physical handle generated by the first AI accelerator hardware;
• returning, to the AI workload, a virtual handle in place of the physical handle, wherein the virtual handle is stored in a host-side data structure for use in a subsequent GPU-related API call;
• maintaining, by the checkpointing system, a mapping between the virtual handle and the physical handle generated by the first AI accelerator hardware;
• monitoring performance of the first AI workload;
• and
• based on the monitoring, migrating the first AI workload to a second resource that is distinct from the first resource and comprising second AI accelerator hardware, the migrating comprising;
• checkpointing the AI workload by saving a checkpoint state comprising the log and the mapping;
• restoring the GPU-related API call on the second resource;
• generating a new physical handle corresponding to the GPU object;
• updating the mapping to associate the virtual handle with the new physical handle;
• and
• resuming execution of the AI workload using the virtual handle.

A computer storage medium having computer-executable instructions for migrating an artificial intelligence (AI) workload wherein, upon execution by a processor, the computer-executable instructions cause the processor to perform operations comprising:

• executing the AI workload on a first resource comprising a graphics processor unit (GPU) for AI acceleration, the executing comprising:
• executing a GPU-related application programming interface (API) call that generates a physical handle to a GPU object, the GPU object comprising at least one of a stream, an event, or a library context;
• storing the GPU-related API call in a log;
• intercepting, by a checkpointing system, the physical handle generated by the GPU;
• returning, to the AI workload, a virtual handle in place of the physical handle, wherein the virtual handle is stored in host-side data structure for use in a subsequent GPU-related API call;
• maintaining, by the checkpointing system, a mapping between the virtual handle and the physical handle generated by the GPU;
• checkpointing the AI workload by saving a checkpoint state comprising the log and the mapping;
• migrating the AI workload to a second resource that is separate from the first resource;
• restoring the checkpoint state of the AI workload on the second resource, the restoring comprising:
• replaying the log, including the GPU-related API call, on the second resource;
• generating a new physical handle corresponding to the GPU object;
• updating the mapping to associate the virtual handle with the new physical handle;
• and
• resuming execution of the AI workload using the virtual handle.