truss 0.11.2rc503__py3-none-any.whl → 0.11.2rc504__py3-none-any.whl

{truss-0.11.2rc503.dist-info → truss-0.11.2rc504.dist-info}/RECORD

@@ -2,29 +2,29 @@ truss/__init__.py,sha256=CoUcP6vx_pocyemRmpbCPlndkHhdMkABAlr0ZXVuPCk,1163
 truss/api/__init__.py,sha256=spBAa_m1pItiid97iDLKPmumgAkSirPkv-E8RWMZyOk,5090
 truss/api/definitions.py,sha256=QAaIBqL59Q-R7HtLcXcoeCIWBN2HqOzApdFX0PpCq2s,1604
 truss/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truss/base/constants.py,sha256=qwNNkd9EOAuiTxYLVccJaiPCNRayBAFvyj_GisYOT3I,3488
+truss/base/constants.py,sha256=sExArdnuGg83z83XMgaQ4b8SS3V_j_bJEpOATDGJzpE,3600
 truss/base/custom_types.py,sha256=FUSIT2lPOQb6gfg6IzT63YBV8r8L6NIZ0D74Fp3e_jQ,2835
 truss/base/errors.py,sha256=zDVLEvseTChdPP0oNhBBQCtQUtZJUaof5zeWMIjqz6o,691
 truss/base/trt_llm_config.py,sha256=CRz3AqGDAyv8YpcBWXUrnfjvNAauyo3yf8ZOGVsSt6g,32782
 truss/base/truss_config.py,sha256=7CtiJIwMHtDU8Wzn8UTJUVVunD0pWFl4QUVycK2aIpY,28055
 truss/base/truss_spec.py,sha256=jFVF79CXoEEspl2kXBAPyi-rwISReIGTdobGpaIhwJw,5979
-truss/cli/chains_commands.py,sha256=bqOXQ-0RPS66vSP_OPQdJ5dvctGiVrsGoSUMbURGdSI,16970
+truss/cli/chains_commands.py,sha256=Kpa5mCg6URAJQE2ZmZfVQFhjBHEitKT28tKiW0H6XAI,17406
 truss/cli/cli.py,sha256=PaMkuwXZflkU7sa1tEoT_Zmy-iBkEZs1m4IVqcieaeo,30367
 truss/cli/remote_cli.py,sha256=G_xCKRXzgkCmkiZJhUFfsv5YSVgde1jLA5LPQitpZgI,1905
-truss/cli/train_commands.py,sha256=GDye7yXGL_nQvXAlY5MWsdj5x0zYOvcQw0Ubn14TiRU,14365
+truss/cli/train_commands.py,sha256=TZhtvofviWQF34pYppRCaQ6qayTsvPnx6afTrYbFpOM,17319
 truss/cli/logs/base_watcher.py,sha256=KKyd7lIrdaEeDVt8EtjMioSPGVpLyOcF0ewyzE_GGdQ,2785
 truss/cli/logs/model_log_watcher.py,sha256=NACcP-wkcaroYa2Cb9BZC7Yr0554WZa_FSM2LXOf4A8,1263
 truss/cli/logs/training_log_watcher.py,sha256=r6HRqrLnz-PiKTUXiDYYxg4ZnP8vYcXlEX1YmgHhzlo,1173
 truss/cli/logs/utils.py,sha256=z-U_FG4BUzdZLbE3BnXb4DZQ0zt3LSZ3PiQpLaDuc3o,1031
 truss/cli/train/common.py,sha256=xTR41U5FeSndXfNBBHF9wF5XwZH1sOIVFlv-XHjsKIU,1547
-truss/cli/train/core.py,sha256=dAmetxKqSc4bQPnVS_8WLfNsw1L7vLT2tU02BVwRPgc,20206
+truss/cli/train/core.py,sha256=4vPnREmaJh8R_rlwR0_H5NRaXhdyY2g07w11uab-9qw,25908
 truss/cli/train/deploy_from_checkpoint_config.yml,sha256=mktaVrfhN8Kjx1UveC4xr-gTW-kjwbHvq6bx_LpO-Wg,371
 truss/cli/train/deploy_from_checkpoint_config_whisper.yml,sha256=6GbOorYC8ml0UyOUvuBpFO_fuYtYE646JqsalR-D4oY,406
 truss/cli/train/metrics_watcher.py,sha256=smz-zrEsBj_-wJHI0pAZ-EAPrvfCWzq1eQjGiFNM-Mk,12755
 truss/cli/train/poller.py,sha256=TGRzELxsicga0bEXewSX1ujw6lfPmDnHd6nr8zvOFO8,3550
 truss/cli/train/types.py,sha256=alGtr4Q71GeB65PpGMhsoKygw4k_ncR6MKIP1ioP8rI,951
 truss/cli/train/deploy_checkpoints/__init__.py,sha256=wL-M2yu8PxO2tFvjwshXAfPnB-5TlvsBp2v_bdzimRU,99
-truss/cli/train/deploy_checkpoints/deploy_checkpoints.py,sha256=xfblHi3py7GDgY24NcuAaDKzcQeOm67rjtWOK6vAEe4,17352
+truss/cli/train/deploy_checkpoints/deploy_checkpoints.py,sha256=nJHUjR4f4_13mFPNetWKq7ecsqr-cradbv3RBBDj2pk,17364
 truss/cli/train/deploy_checkpoints/deploy_checkpoints_helpers.py,sha256=6x5nS_HnWYtS9vi-Pg8akzrJk9L_agjvFhm5EFh1m6Y,1964
 truss/cli/train/deploy_checkpoints/deploy_full_checkpoints.py,sha256=FYRG5KTMlxEMZS-RA_m2gp1wuqWbSpqt2RhdQfLibhA,3968
 truss/cli/train/deploy_checkpoints/deploy_lora_checkpoints.py,sha256=P91dIAzuhl2GlzmrWwCcYI7uCMT1Lm7C79JQHM_exN4,4442
@@ -66,15 +66,15 @@ truss/remote/baseten/utils/time.py,sha256=Ry9GMjYnbIGYVIGwtmv4V8ljWjvdcaCf5NOQzl
 truss/remote/baseten/utils/transfer.py,sha256=d3VptuQb6M1nyS6kz0BAfeOYDLkMKUjatJXpY-mp-As,1548
 truss/templates/README.md.jinja,sha256=N7CJdyldZuJamj5jLh47le0hFBdu9irVsTBqoxhPNPQ,2476
 truss/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truss/templates/base.Dockerfile.jinja,sha256=3deYATR6gbgNmmYhg4yXopbOUnICSbXCfjeGU8mdUaQ,5734
+truss/templates/base.Dockerfile.jinja,sha256=irked6fWbiZ4tMkhR3zi3njpaaI9bANVqq7PTjp_Tmc,5610
 truss/templates/cache.Dockerfile.jinja,sha256=1qZqDo1phrcqi-Vwol-VafYJkADsBbQWU6huQ-_1x00,1146
 truss/templates/cache_requirements.txt,sha256=xoPoJ-OVnf1z6oq_RVM3vCr3ionByyqMLj7wGs61nUs,87
 truss/templates/copy_cache_files.Dockerfile.jinja,sha256=Os5zFdYLZ_AfCRGq4RcpVTObOTwL7zvmwYcvOzd_Zqo,126
 truss/templates/docker_server_requirements.txt,sha256=PyhOPKAmKW1N2vLvTfLMwsEtuGpoRrbWuNo7tT6v2Mc,18
-truss/templates/server.Dockerfile.jinja,sha256=fNrCi1sGefjGlu2JOzLu8E7PiwZvpoq8JgW_BBhvya0,7219
-truss/templates/control/requirements.txt,sha256=Kk0tYID7trPk5gwX38Wrt2-YGWZAXFJCJRcqJ8ZzCjc,251
+truss/templates/server.Dockerfile.jinja,sha256=CUYnF_hgxPGq2re7__0UPWlwzOHMoFkxp6NVKi3U16s,7071
+truss/templates/control/requirements.txt,sha256=tJGr83WoE0CZm2FrloZ9VScK84q-_FTuVXjDYrexhW0,250
 truss/templates/control/control/application.py,sha256=jYeta6hWe1SkfLL3W4IDmdYjg3ZuKqI_UagWYs5RB_E,3793
-truss/templates/control/control/endpoints.py,sha256=FM-sgao7I3gMoUTasM3Xq_g2LDoJQe75JxIoaQxzeNo,10031
+truss/templates/control/control/endpoints.py,sha256=VQ1lvZjFvR091yRkiFdvXw1Q7PiNGXT9rJwY7_sX6yg,11828
 truss/templates/control/control/server.py,sha256=R4Y219i1dcz0kkksN8obLoX-YXWGo9iW1igindyG50c,3128
 truss/templates/control/control/helpers/context_managers.py,sha256=W6dyFgLBhPa5meqrOb3w_phMtKfaJI-GhwUfpiycDc8,413
 truss/templates/control/control/helpers/custom_types.py,sha256=n_lTudtLTpy4oPV3aDdJ4X2rh3KCV5btYO9UnTeUouQ,5471
@@ -84,7 +84,7 @@ truss/templates/control/control/helpers/inference_server_process_controller.py,s
 truss/templates/control/control/helpers/inference_server_starter.py,sha256=Fz2Puijro6Cc5cvTsAqOveNJbBQR_ARYJXl4lvETJ8Y,2633
 truss/templates/control/control/helpers/truss_patch/__init__.py,sha256=CXZdUV_ylqLTJrKuFpvSnUT6PUFrZrMF2y6jiHbdaKU,998
 truss/templates/control/control/helpers/truss_patch/model_code_patch_applier.py,sha256=LTIIADLz0wRn7V49j64dU1U7Hbta9YLde3pb5YZWvzQ,2001
-truss/templates/control/control/helpers/truss_patch/model_container_patch_applier.py,sha256=62TDVaDmgAH0-X116xSDnNTOFEgUQH4sNJr0aALFl_0,7149
+truss/templates/control/control/helpers/truss_patch/model_container_patch_applier.py,sha256=rxJzXwNmrr_2FOM9gZSE3TDbV-Bpwz67fv1dcMsQ9EE,6420
 truss/templates/control/control/helpers/truss_patch/requirement_name_identifier.py,sha256=CL3KEAj4B3ApMQShd7TI5umXVbazLZY5StrNlwHwWtc,1995
 truss/templates/control/control/helpers/truss_patch/system_packages.py,sha256=IYh1CVU_kooAvtSGXKQDDWnNdOhlv7ENWagsL1wvhgw,208
 truss/templates/custom/examples.yaml,sha256=2UcCtEdavImWmiCtj31ckBlAKVOwNMC5AwMIIznKDag,48
@@ -94,13 +94,12 @@ truss/templates/custom_python_dx/my_model.py,sha256=NG75mQ6wxzB1BYUemDFZvRLBET-U
 truss/templates/docker_server/proxy.conf.jinja,sha256=Lg-PcZzKflG85exZKHNgW_I6r0mATV8AtOIBaE40-RM,1669
 truss/templates/docker_server/supervisord.conf.jinja,sha256=dd37fwZE--cutrvOUCqEyJQQQhlp61H2IUs2huKWsSk,1808
 truss/templates/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truss/templates/server/entrypoint.sh,sha256=aRqtkHt_gHRgEa2S1sGnDoaFaZtA43NbvKYUy5CtzcM,881
 truss/templates/server/main.py,sha256=kWXrdD8z8IpamyWxc8qcvd5ck9gM1Kz2QH5qHJCnmOQ,222
 truss/templates/server/model_wrapper.py,sha256=k75VVISwwlsx5EGb82UZsu8kCM_i6Yi3-Hd0-Kpm1yo,42055
 truss/templates/server/requirements.txt,sha256=XblmpfxAmRo3X1V_9oMj8yjdpZ5Wk-C2oa3z6nq4OGw,672
-truss/templates/server/truss_server.py,sha256=ob_nceeGtFPZzKKdk_ZZGLoZrJOGE6hR52xM1sPR97A,19498
+truss/templates/server/truss_server.py,sha256=noXfGJMsKIhgF4oI_8LC1UHkcx8Vg8nGSITZJ_bkRFQ,19598
 truss/templates/server/common/__init__.py,sha256=qHIqr68L5Tn4mV6S-PbORpcuJ4jmtBR8aCuRTIWDvNo,85
-truss/templates/server/common/errors.py,sha256=qWeZlmNI8ZGbZbOIp_mtS6IKvUFIzhj3QH8zp-xTp9o,8554
+truss/templates/server/common/errors.py,sha256=My0P6-Y7imVTICIhazHT0vlSu3XJDH7As06OyVzu4Do,8589
 truss/templates/server/common/patches.py,sha256=uEOzvDnXsHOkTSa8zygGYuR4GHhrFNVHNQc5peJcwvo,1393
 truss/templates/server/common/retry.py,sha256=dtz6yvwLoY0i55FnxECz57zEOKjAhGMYvvM-k9jiR9c,624
 truss/templates/server/common/schema.py,sha256=WLFtVyEKmk4whg5_gk6Gt1vOD6wM5fWKLb4zNuD0bkw,6042
@@ -113,6 +112,8 @@ truss/templates/shared/log_config.py,sha256=l9udyu4VKHZePlfK9LQEd5TOUUodPuehypsX
 truss/templates/shared/secrets_resolver.py,sha256=3prDe3Q06NTmUEe7KCW-W4TD1CzGck9lpDG789209z4,2110
 truss/templates/shared/serialization.py,sha256=_WC_2PPkRi-MdTwxwjG8LKQptnHi4sANfpOlKWevqWc,3736
 truss/templates/shared/util.py,sha256=dPgFF4iL_YkeC6Kf8tZUHJH60rbpskHwVPh0ONLGaQM,2222
+truss/templates/train/config.py,sha256=aQJ3lsyVRlq6edjjZq4_Anz1bZVwkjLdclmZPJTdo1k,1626
+truss/templates/train/run.sh,sha256=2rimigJOn6yg4DguRfOJWkzm77X-meNSYXnidLafqNg,346
 truss/templates/trtllm-audio/model/model.py,sha256=o38QqW57b1lf8O_td1lW_AojZZ8R_qAZCgzOWtoIse8,1619
 truss/templates/trtllm-audio/packages/sigint_patch.py,sha256=t6pYpVwgQsLCgcxQq7-V3scr9ZOiIxtYSpy9LCfdNTk,414
 truss/templates/trtllm-audio/packages/whisper_trt/__init__.py,sha256=5ZQfVlwtkWrnjYiuBIVSviYDhV-kksygDkHEWBS_ijM,7065
@@ -141,8 +142,9 @@ truss/tests/test_truss_handle.py,sha256=-xz9VXkecXDTslmQZ-dmUmQLnvD0uumRqHS2uvGl
 truss/tests/test_util.py,sha256=hs1bNMkXKEdoPRx4Nw-NAEdoibR92OubZuADGmbiYsQ,1344
 truss/tests/cli/test_cli.py,sha256=yfbVS5u1hnAmmA8mJ539vj3lhH-JVGUvC4Q_Mbort44,787
 truss/tests/cli/train/test_cache_view.py,sha256=aVRCh3atRpFbJqyYgq7N-vAW0DiKMftQ7ajUqO2ClOg,22606
-truss/tests/cli/train/test_deploy_checkpoints.py,sha256=wQZ3DPLPAyXE3iaQiyHJTBO15v_gXN44eDk1StYkKmM,44764
+truss/tests/cli/train/test_deploy_checkpoints.py,sha256=lDk88uAUPYatJ30JKVVtJDdXv_zWNk1nxXFyUH6IVGw,44800
 truss/tests/cli/train/test_train_cli_core.py,sha256=vzYfxKdwoa3NaFMrVZbSg5qOoLXivMvZXN1ClQirGTQ,16148
+truss/tests/cli/train/test_train_init.py,sha256=pv8BfyLlVG0QtdowTziITjKa_OE1KigatmAGx8XSZrM,17238
 truss/tests/cli/train/resources/test_deploy_from_checkpoint_config.yml,sha256=GF7r9l0KaeXiUYCPSBpeMPd2QG6PeWWyI12NdbqLOgc,1930
 truss/tests/contexts/image_builder/test_serving_image_builder.py,sha256=16niCXZnuxFHXYQw2vPFZ8svSZafkH5DT0Gx3Z9Xdd8,22377
 truss/tests/contexts/local_loader/test_load_local.py,sha256=D1qMH2IpYA2j5009v50QMgUnKdeOsX15ndkwXe10a4E,801
@@ -162,7 +164,7 @@ truss/tests/remote/baseten/test_auth.py,sha256=ttu4bDnmwGfo3oiNut4HVGnh-QnjAefwZ
 truss/tests/remote/baseten/test_core.py,sha256=6NzJTDmoSUv6Muy1LFEYIUg10-cqw-hbLyeTSWcdNjY,26117
 truss/tests/remote/baseten/test_remote.py,sha256=y1qSPL1t7dBeYI3xMFn436fttG7wkYdAoENTz7qKObg,23634
 truss/tests/remote/baseten/test_service.py,sha256=ufZbtQlBNIzFCxRt_iE-APLpWbVw_3ViUpSh6H9W5nU,1945
-truss/tests/templates/control/control/test_endpoints.py,sha256=tGU3w8zOKC8LfWGdhp-TlV7E603KXg2xGwpqDdf8Pnw,3385
+truss/tests/templates/control/control/test_endpoints.py,sha256=fxTiiCR0ltaHCL_-v-22Ie1qVgnch1lqcj3w0U3R-fk,3644
 truss/tests/templates/control/control/test_server.py,sha256=r1O3VEK9eoIL2-cg8nYLXYct_H3jf5rGp1wLT1KBdeA,9488
 truss/tests/templates/control/control/test_server_integration.py,sha256=EdDY3nLzjrRCJ5LI5yZsNCEImSRkxTL7Rn9mGnK67zA,11837
 truss/tests/templates/control/control/helpers/test_context_managers.py,sha256=3LoonRaKu_UvhaWs1eNmEQCZq-iJ3aIjI0Mn4amC8Bw,283
@@ -346,27 +348,28 @@ truss_chains/__init__.py,sha256=QDw1YwdqMaQpz5Oltu2Eq2vzEX9fDrMoqnhtbeh60i4,1278
 truss_chains/framework.py,sha256=CS7tSegPe2Q8UUT6CDkrtSrB3utr_1QN1jTEPjrj5Ug,67519
 truss_chains/private_types.py,sha256=6CaQEPawFLXjEbJ-01lqfexJtUIekF_q61LNENWegFo,8917
 truss_chains/public_api.py,sha256=0AXV6UdZIFAMycUNG_klgo4aLFmBZeKGfrulZEWzR0M,9532
-truss_chains/public_types.py,sha256=q8Oet6MpECW1FhWW25SCExpZhmk4cFmEsqrO30oZIMw,29112
+truss_chains/public_types.py,sha256=Am1vc5pWuEDs65UQp8Be4iOU05kMXjPG3QrJfNXDmHs,29225
 truss_chains/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/pydantic_numpy.py,sha256=MG8Ji_Inwo_JSfM2n7TPj8B-nbrBlDYsY3SOeBwD8fE,4289
 truss_chains/streaming.py,sha256=DGl2LEAN67YwP7Nn9MK488KmYc4KopWmcHuE6WjyO1Q,12521
 truss_chains/utils.py,sha256=LvpCG2lnN6dqPqyX3PwLH9tyjUzqQN3N4WeEFROMHak,6291
 truss_chains/deployment/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truss_chains/deployment/code_gen.py,sha256=AmAUZ3h1hP3uYkl3J6o096K5RFLuBOP7kOFSnFC_C4U,32568
-truss_chains/deployment/deployment_client.py,sha256=haFiVmQek42ewlN_YflBaRDQT4ZYbmT20tvvJOkcUX0,32899
+truss_chains/deployment/code_gen.py,sha256=IBOYdhsWUyW_sBVhlEQAhvwxKcsGflDjgmR-1HyJJLg,32666
+truss_chains/deployment/deployment_client.py,sha256=2paNyBjrpFTxROP0YrmJMUlH6o8mrkF-iPms7VJhLdA,34017
 truss_chains/reference_code/reference_chainlet.py,sha256=5feSeqGtrHDbldkfZCfX2R5YbbW0Uhc35mhaP2pXrHw,1340
 truss_chains/reference_code/reference_model.py,sha256=emH3hb23E_nbP98I37PGp1Xk1hz3g3lQ00tiLo55cSM,322
 truss_chains/remote_chainlet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/remote_chainlet/model_skeleton.py,sha256=8ZReLOO2MLcdg7bNZ61C-6j-e68i2Z-fFlyV3sz0qH8,2376
 truss_chains/remote_chainlet/stub.py,sha256=Y2gDUzMY9WRaQNHIz-o4dfLUfFyYV9dUhIRQcfgrY8g,17209
-truss_chains/remote_chainlet/utils.py,sha256=O_5P-VAUvg0cegEW1uKCOf5EBwD8rEGYVoGMivOmc7k,22374
-truss_train/__init__.py,sha256=7hE6j6-u6UGzCGaNp3CsCN0kAVjBus1Ekups-Bk0fi4,837
-truss_train/definitions.py,sha256=V985HhY4rdXL10DZxpFEpze9ScxzWErMht4WwaPknGU,6789
+truss_chains/remote_chainlet/utils.py,sha256=Zn3GZRvK8f65WUa-qa-8uPFZ2pD7ukRFxbLOvT-BL0Q,24063
+truss_train/__init__.py,sha256=A3MzRPMInZfmzLvPpZI7gdKgshAVCw6bwhU-6JYU2zs,939
+truss_train/definitions.py,sha256=3wVxkxMtHlcc-hb2umtj74FjA9TjenfiPTX7EQSh6zw,8245
 truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,3074
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
-truss-0.11.2rc503.dist-info/METADATA,sha256=ibPdGGEmpNg-1bG8GcmjRw0UR1KB6Dak-4uHszZ9_qo,6674
-truss-0.11.2rc503.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.2rc503.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.2rc503.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.2rc503.dist-info/RECORD,,
+truss_train/restore_from_checkpoint.py,sha256=KmJuTUVpvtvlkEClcmllxAF2TKgbp-FuzfblfGh06XA,1239
+truss-0.11.2rc504.dist-info/METADATA,sha256=kmeypQ9CmBsHrPE-nB0Xw7DYTGCirhsQIW6wbP9yiSU,6674
+truss-0.11.2rc504.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.2rc504.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.2rc504.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.2rc504.dist-info/RECORD,,

truss_chains/deployment/code_gen.py

@@ -544,7 +544,11 @@ async def websocket(self, websocket: fastapi.WebSocket) -> None:
         )"""
     return _Source(
         src=src,
-        imports={"import fastapi", "from truss_chains.remote_chainlet import utils"},
+        imports={
+            "import fastapi",
+            "from starlette.websockets import WebSocketState",
+            "from truss_chains.remote_chainlet import utils",
+        },
     )
 
 

truss_chains/deployment/deployment_client.py

@@ -138,7 +138,12 @@ class ChainService(abc.ABC):
 
 def _generate_chainlet_artifacts(
     options: private_types.PushOptions, entrypoint: Type[private_types.ABCChainlet]
-) -> tuple[b10_types.ChainletArtifact, list[b10_types.ChainletArtifact], bool]:
+) -> tuple[
+    b10_types.ChainletArtifact,
+    list[b10_types.ChainletArtifact],
+    bool,
+    Optional[private_types.ChainletAPIDescriptor],
+]:
     chain_root = _get_chain_root(entrypoint)
     entrypoint_artifact: Optional[b10_types.ChainletArtifact] = None
     dependency_artifacts: list[b10_types.ChainletArtifact] = []
@@ -192,7 +197,19 @@ def _generate_chainlet_artifacts(
 
     assert entrypoint_artifact is not None
 
-    return entrypoint_artifact, dependency_artifacts, has_engine_builder_chainlets
+    # Find the entrypoint descriptor
+    entrypoint_descriptor = None
+    for chainlet_descriptor in _get_ordered_dependencies([entrypoint]):
+        if chainlet_descriptor.chainlet_cls == entrypoint:
+            entrypoint_descriptor = chainlet_descriptor
+            break
+
+    return (
+        entrypoint_artifact,
+        dependency_artifacts,
+        has_engine_builder_chainlets,
+        entrypoint_descriptor,
+    )
 
 
 @framework.raise_validation_errors_before
@@ -201,9 +218,12 @@ def push(
     options: private_types.PushOptions,
     progress_bar: Optional[Type["progress.Progress"]] = None,
 ) -> Optional[ChainService]:
-    entrypoint_artifact, dependency_artifacts, has_engine_builder_chainlets = (
-        _generate_chainlet_artifacts(options, entrypoint)
-    )
+    (
+        entrypoint_artifact,
+        dependency_artifacts,
+        has_engine_builder_chainlets,
+        entrypoint_descriptor,
+    ) = _generate_chainlet_artifacts(options, entrypoint)
     if options.only_generate_trusses:
         return None
     if isinstance(options, private_types.PushOptionsBaseten):
@@ -213,7 +233,11 @@ def push(
                 "not supportd, push with `--publish`."
             )
         return _create_baseten_chain(
-            options, entrypoint_artifact, dependency_artifacts, progress_bar
+            options,
+            entrypoint_artifact,
+            dependency_artifacts,
+            progress_bar,
+            entrypoint_descriptor,
         )
     elif isinstance(options, private_types.PushOptionsLocalDocker):
         if has_engine_builder_chainlets:
@@ -369,16 +393,19 @@ def _create_docker_chain(
 class BasetenChainService(ChainService):
     _chain_deployment_handle: b10_core.ChainDeploymentHandleAtomic
     _remote: b10_remote.BasetenRemote
+    _entrypoint_descriptor: Optional[private_types.ChainletAPIDescriptor]
 
     def __init__(
         self,
         name: str,
         chain_deployment_handle: b10_core.ChainDeploymentHandleAtomic,
         remote: b10_remote.BasetenRemote,
+        entrypoint_descriptor: Optional[private_types.ChainletAPIDescriptor] = None,
     ) -> None:
         super().__init__(name)
         self._chain_deployment_handle = chain_deployment_handle
         self._remote = remote
+        self._entrypoint_descriptor = entrypoint_descriptor
 
     @property
     def run_remote_url(self) -> str:
@@ -393,6 +420,13 @@ class BasetenChainService(ChainService):
             is_draft=handle.is_draft,
         )
 
+    @property
+    def is_websocket(self) -> bool:
+        """Check if the entrypoint uses websockets."""
+        if self._entrypoint_descriptor is None:
+            return False
+        return self._entrypoint_descriptor.endpoint.is_websocket
+
     def run_remote(self, json_data: Dict) -> Any:
         """Invokes the entrypoint with JSON data.
 
@@ -462,6 +496,7 @@ def _create_baseten_chain(
     entrypoint_artifact: b10_types.ChainletArtifact,
     dependency_artifacts: list[b10_types.ChainletArtifact],
     progress_bar: Optional[Type["progress.Progress"]],
+    entrypoint_descriptor: Optional[private_types.ChainletAPIDescriptor] = None,
 ):
     logging.info(
         f"Pushing Chain '{baseten_options.chain_name}' to Baseten "
@@ -491,7 +526,10 @@ def _create_baseten_chain(
         progress_bar=progress_bar,
     )
     return BasetenChainService(
-        baseten_options.chain_name, chain_deployment_handle, remote_provider
+        baseten_options.chain_name,
+        chain_deployment_handle,
+        remote_provider,
+        entrypoint_descriptor,
     )
 
 

truss_chains/public_types.py

@@ -473,6 +473,7 @@ class WebSocketProtocol(Protocol):
 
     async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...
 
+    async def receive(self) -> Union[str, bytes]: ...
     async def receive_text(self) -> str: ...
     async def receive_bytes(self) -> bytes: ...
     async def receive_json(self) -> Any: ...
@@ -481,9 +482,11 @@ class WebSocketProtocol(Protocol):
     async def send_bytes(self, data: bytes) -> None: ...
     async def send_json(self, data: Any) -> None: ...
 
-    def iter_text(self) -> AsyncIterator[str]: ...
-    def iter_bytes(self) -> AsyncIterator[bytes]: ...
-    def iter_json(self) -> AsyncIterator[Any]: ...
+    async def iter_text(self) -> AsyncIterator[str]: ...
+    async def iter_bytes(self) -> AsyncIterator[bytes]: ...
+    async def iter_json(self) -> AsyncIterator[Any]: ...
+
+    def is_connected(self) -> bool: ...
 
 
 class EngineBuilderLLMInput(pydantic.BaseModel):

truss_chains/remote_chainlet/utils.py

@@ -11,6 +11,7 @@ import textwrap
 import threading
 import time
 import traceback
+import typing
 from collections.abc import AsyncIterator
 from typing import (
     TYPE_CHECKING,
@@ -38,7 +39,7 @@ if TYPE_CHECKING:
 try:
     import prometheus_client
 except ImportError:
-    logging.warning("Optional `prometheus_client` is not installed. ")
+    logging.warning("Optional `prometheus_client` is not installed.")
 
     class _NoOpMetric:
         def labels(self, *args: object, **kwargs: object) -> "_NoOpMetric":
@@ -64,6 +65,26 @@ except ImportError:
             return _NoOpMetric()
 
 
+try:
+    from fastapi import WebSocketDisconnect
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketDisconnect, in case local environment doesn't have
+    # fastapi.
+    class WebSocketDisconnect(Exception):  # type: ignore[no-redef]
+        def __init__(self, code: int, reason: Optional[str] = None):
+            super().__init__()
+            self.code = code
+            self.reason = reason
+
+
+try:
+    from starlette.websockets import WebSocketState
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketState, in case local environment doesn't have starlette.
+    class WebSocketState:  # type: ignore[no-redef]
+        CONNECTED = "connected"
+
+
 T = TypeVar("T")
 
 _LockT = TypeVar("_LockT", bound=Union[threading.Lock, asyncio.Lock])
@@ -586,6 +607,18 @@ class WebsocketWrapperFastAPI:
     async def close(self, code: int = 1000, reason: Optional[str] = None) -> None:
         await self._websocket.close(code=code, reason=reason)
 
+    async def receive(self) -> Union[str, bytes]:
+        message = await self._websocket.receive()
+
+        if message.get("type") == "websocket.disconnect":
+            # NB(nikhil): Mimics FastAPI `_raise_on_disconnect`, since otherwise the user has no
+            # way of detecting that the client disconnected.
+            raise WebSocketDisconnect(message["code"], message.get("reason"))
+        elif message.get("text"):
+            return typing.cast(str, message["text"])
+        else:
+            return typing.cast(bytes, message["bytes"])
+
     async def receive_text(self) -> str:
         return await self._websocket.receive_text()
 
@@ -605,13 +638,19 @@ class WebsocketWrapperFastAPI:
         await self._websocket.send_json(data)
 
     async def iter_text(self) -> AsyncIterator[str]:
-        while True:
-            yield await self.receive_text()
+        return self._websocket.iter_text()
 
     async def iter_bytes(self) -> AsyncIterator[bytes]:
-        while True:
-            yield await self.receive_bytes()
+        return self._websocket.iter_bytes()
 
     async def iter_json(self) -> AsyncIterator[Any]:
-        while True:
-            yield await self.receive_json()
+        return self._websocket.iter_json()
+
+    async def is_connected(self) -> bool:
+        # NB(nikhil): This isn't a foolproof mechanism for detecting whether a websocket
+        # connection is actually alive, ping/pong messages are best suited for that. However,
+        # as a heuristic to determine if a message is safe to send, this can do a pretty good job.
+        return (
+            self._websocket.application_state == WebSocketState.CONNECTED
+            and self._websocket.client_state == WebSocketState.CONNECTED
+        )

truss_train/__init__.py

@@ -1,5 +1,6 @@
 from truss_train.definitions import (
     AWSIAMDockerAuth,
+    BasetenCheckpoint,
     CacheConfig,
     CheckpointingConfig,
     CheckpointList,
@@ -10,6 +11,7 @@ from truss_train.definitions import (
     FullCheckpoint,
     GCPServiceAccountJSONDockerAuth,
     Image,
+    LoadCheckpointConfig,
     LoRACheckpoint,
     LoRADetails,
     ModelWeightsFormat,
@@ -36,6 +38,8 @@ __all__ = [
     "CacheConfig",
     "AWSIAMDockerAuth",
     "GCPServiceAccountJSONDockerAuth",
+    "LoadCheckpointConfig",
+    "BasetenCheckpoint",
     "DockerAuth",
     "Image",
 ]

truss_train/definitions.py

@@ -1,11 +1,11 @@
 import enum
 from abc import ABC
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 import pydantic
 from pydantic import field_validator, model_validator
 
-from truss.base import custom_types, truss_config
+from truss.base import constants, custom_types, truss_config
 
 DEFAULT_LORA_RANK = 16
 
@@ -56,6 +56,50 @@ class Compute(custom_types.SafeModelNoExtra):
         )
 
 
+class _CheckpointBase(custom_types.SafeModelNoExtra):
+    typ: str
+
+
+class _BasetenLatestCheckpoint(_CheckpointBase):
+    job_id: Optional[str] = None
+    project_name: Optional[str] = None
+    typ: Literal["baseten_latest_checkpoint"] = "baseten_latest_checkpoint"
+
+
+class _BasetenNamedCheckpoint(_CheckpointBase):
+    checkpoint_name: str
+    job_id: Optional[str]
+    project_name: Optional[str]
+    typ: Literal["baseten_named_checkpoint"] = "baseten_named_checkpoint"
+
+
+class BasetenCheckpoint:
+    @staticmethod
+    def from_latest_checkpoint(
+        project_name: Optional[str] = None, job_id: Optional[str] = None
+    ) -> _BasetenLatestCheckpoint:
+        return _BasetenLatestCheckpoint(project_name=project_name, job_id=job_id)
+
+    @classmethod
+    def from_named_checkpoint(
+        cls,
+        checkpoint_name: str,
+        project_name: Optional[str] = None,
+        job_id: Optional[str] = None,
+    ) -> _BasetenNamedCheckpoint:
+        return _BasetenNamedCheckpoint(
+            checkpoint_name=checkpoint_name, project_name=project_name, job_id=job_id
+        )
+
+
+class LoadCheckpointConfig(custom_types.SafeModelNoExtra):
+    enabled: bool = False
+    checkpoints: List[Union[_BasetenLatestCheckpoint, _BasetenNamedCheckpoint]] = [
+        _BasetenLatestCheckpoint()
+    ]
+    download_folder: str = constants.DEFAULT_TRAINING_CHECKPOINT_FOLDER
+
+
 class CheckpointingConfig(custom_types.SafeModelNoExtra):
     enabled: bool = False
     checkpoint_path: Optional[str] = None
@@ -72,6 +116,7 @@ class Runtime(custom_types.SafeModelNoExtra):
     environment_variables: Dict[str, Union[str, SecretReference]] = {}
     enable_cache: Optional[bool] = None
     checkpointing_config: CheckpointingConfig = CheckpointingConfig()
+    load_checkpoint_config: Optional[LoadCheckpointConfig] = None
     cache_config: Optional[CacheConfig] = None
 
     @model_validator(mode="before")

truss_train/restore_from_checkpoint.py

@@ -0,0 +1,42 @@
+from truss_train import (
+    BasetenCheckpoint,
+    CheckpointingConfig,
+    Image,
+    LoadCheckpointConfig,
+    Runtime,
+    TrainingJob,
+    TrainingProject,
+)
+
+load_checkpoint_config = LoadCheckpointConfig(enabled=True)
+
+load_from_most_recent_checkpoint = BasetenCheckpoint.from_latest_checkpoint()
+
+load_most_recent_checkpoint = BasetenCheckpoint.from_latest_checkpoint(
+    job_id="lqz4pw4",  # Optional
+    project_name="first-project",  # Optional
+)
+
+load_from_named_checkpoint = BasetenCheckpoint.from_named_checkpoint(
+    checkpoint_name="checkpoint-24",
+    project_name="first-project",  # Optional
+    job_id="lqz4pw4",  # Optional
+)
+
+load_checkpoint_config = LoadCheckpointConfig(
+    enabled=True,
+    download_folder="/tmp/custom_location",  # default is None -> default path set by server-side
+    checkpoints=[load_from_most_recent_checkpoint, load_from_named_checkpoint],
+)
+
+checkpointing_config = CheckpointingConfig(enabled=True)
+
+job = TrainingJob(
+    image=Image(base_image="ghcr.io/baseten-ai/truss-train-base:latest"),
+    runtime=Runtime(
+        checkpointing_config=checkpointing_config,
+        load_checkpoint_config=load_checkpoint_config,
+    ),
+)
+
+project = TrainingProject(name="new-project", job=job)

truss/templates/server/entrypoint.sh

@@ -1,32 +0,0 @@
-#!/bin/bash
-set -e
-
-# Test sudo permissions first
-echo "Testing sudo permissions..."
-sudo -l | grep chown || echo "No chown permissions found in sudo -l"
-
-# Fix ownership of cache directories if they exist
-if [ -d "/cache/model" ]; then
-    echo "Setting ownership for /cache/model"
-    echo "Before: $(ls -ld /cache/model)"
-    if sudo chown -R app:app /cache/model; then
-        echo "chown command succeeded"
-    else
-        echo "chown command failed with exit code: $?"
-    fi
-    echo "After: $(ls -ld /cache/model)"
-fi
-
-if [ -d "/cache/org" ]; then
-    echo "Setting ownership for /cache/org"
-    echo "Before: $(ls -ld /cache/org)"
-    if sudo chown -R app:app /cache/org; then
-        echo "chown command succeeded"
-    else
-        echo "chown command failed with exit code: $?"
-    fi
-    echo "After: $(ls -ld /cache/org)"
-fi
-
-# Execute the original command
-exec "$@"